import numpy as np
import plotly.graph_objects as go
from mrmh_model.lp_model import create_heatmap
from arcgis_terrain import get_terrain_map
import pickle as pkl
from scipy.interpolate import griddata
from scipy import interpolate
from arcgis_terrain import lat_lon2meters
from scipy.io import loadmat


def linear_features(parameters):

    temp = loadmat(parameters['lin_feat_filename'])

    bwlf = temp['BWLF'] # linear features
    bwinac = temp['BWInac'] # inaccessible areas
    szelev = temp['sZelev'] # elevation

    dim = bwlf.shape[0]

    bwlf = np.fliplr(np.rot90(np.rot90(bwlf)))
    bwinac = np.fliplr(np.rot90(np.rot90(bwinac)))

    # rescale z data to match map size
    map_res = parameters['res']
    map_size = [np.int((parameters['xlims'][1] - parameters['xlims'][0])/map_res), np.int((parameters['ylims'][1] - parameters['ylims'][0])/map_res)]
    if map_size[0] != map_size[1]:
        print('not implemented for non-square maps')

    feat_res = 6.6666666666667
    scale_factor = map_res/feat_res # convert scales
    center_idx = [int(dim/2), int(dim/2)] # ics in map coords
    scaled_dim = [np.ceil((map_size[0]/2)*scale_factor), np.ceil((map_size[1]/2)*scale_factor)] # number of 6.7m steps to get to map size

    if scaled_dim[0] >= center_idx[0]:
        print("not fixed. map scales are not right!")

    xcrds_center = np.arange(center_idx[0] - scaled_dim[0], center_idx[0] + scaled_dim[0], dtype=int)
    ycrds_center = np.arange(center_idx[1] - scaled_dim[1], center_idx[1] + scaled_dim[1], dtype=int)

    # trim to correct size:
    lf_idx = np.nonzero(bwlf[xcrds_center,:][:,ycrds_center])
    inac_idx = np.nonzero(bwinac[xcrds_center,:][:,ycrds_center])

    lf_pts = np.array(lf_idx)*feat_res + parameters['xlims'][0] # in units of meters now
    inac_pts = np.array(inac_idx)*feat_res + parameters['xlims'][0]

    # flip y axis
    # lf_pts[1] = parameters['xlims'][1] - lf_pts[1] + parameters['xlims'][0]
    # inac_pts[1] = parameters['xlims'][1] - inac_pts[1] + parameters['xlims'][0]

    # might still need to do some rotations/flips

    return lf_pts, inac_pts

def plot_all(parameters, mc_object, robot_paths, searcher_paths, title = "Sim", smooth_paths = True,
             show_heatmap = True, show_contours = False, cs_name = 'thermal'):

    linear_features(parameters)

    plot_data = []
    color_list = ['blue', 'green', 'cyan', 'magenta', 'lime']
    terrain_data = mc_object.terrain.terrain_data
    heatmap_data = mc_object.p.T

    e = terrain_data[0]
    x = terrain_data[1]
    y = terrain_data[2]

    z_terr = e.T
    x_terr = np.linspace(-(np.max(x)-np.min(x))/2, (np.max(x)-np.min(x))/2, x.shape[0])
    y_terr = np.linspace(-(np.max(y)-np.min(y))/2, (np.max(y)-np.min(y))/2, y.shape[0])

    # heat_map = create_heatmap(z_terr.shape)
    # generate terrain function and use to lift features
    z_terr_func = interpolate.RectBivariateSpline(x_terr, y_terr, z_terr.T)


    for i,path in enumerate(robot_paths):
        z_traj = [val[2] for val in path]
        x_traj = [val[0] for val in path]
        y_traj = [val[1] for val in path]

        z_temp = []
        x_temp = []
        y_temp = []

        # lifting each z point to the terrain
        for j,pt in enumerate(path):
            x_idx = np.argmin(np.abs(x_terr - pt[0]))
            y_idx = np.argmin(np.abs(y_terr - pt[1]))
            if z_terr_func.ev(pt[0], pt[1]) >= (z_traj[j]-4):
                z_traj[j] = z_terr_func.ev(pt[0], pt[1]) + 4
            else:
                pass

            z_val = z_terr_func.ev(pt[0], pt[1]) + 1
            z_temp.append(z_val)
            x_temp.append(x_terr[x_idx])
            y_temp.append(y_terr[y_idx])

        # add some smoothing for long trajectories
        if smooth_paths:
            tck, unused = interpolate.splprep([x_traj, y_traj, z_traj], k=5, s=750)
            unew = np.linspace(0,1,30)
            smoothed_traj = interpolate.splev(unew, tck)
            x_traj_smoothed = smoothed_traj[0]
            y_traj_smoothed = smoothed_traj[1]
            z_traj_smoothed = smoothed_traj[2]
        else:
            x_traj_smoothed = x_traj
            y_traj_smoothed = y_traj
            z_traj_smoothed = z_traj

        # plot_data.append(go.Scatter3d(z = z_temp, x = x_temp, y = y_temp, mode="lines",
        #                               line=dict(color=color_list[i], width=1), name="Proj. {}".format(i+1)))
        plot_data.append(go.Scatter3d(z = z_traj_smoothed, x = x_traj_smoothed, y = y_traj_smoothed,# mode="lines",
                                      line=dict(color=color_list[i], width=4), name="Drone {}".format(i+1)))
        plot_data.append(go.Scatter3d(z = [z_traj_smoothed[0]], x = [x_traj_smoothed[0]], y = [y_traj_smoothed[0]],
                                      mode="markers", showlegend=False,
                                      marker=dict(color=color_list[i], size = 12, line=dict(color='black',width=8)),
                                      name="Drone {}".format(i+1)))

    if show_contours:

        # add surface plot data
        if show_heatmap:
            plot_data.append(go.Surface(z = z_terr, x = x_terr, y = y_terr, colorscale=cs_name,
                                        surfacecolor=heatmap_data, name = "Terrain", showscale=False,
                                        contours = {"z" : {"show" : True, "start": 0, "end": np.max(z_terr), "size": 5}}))
        else:
            plot_data.append(go.Surface(z = z_terr, x = x_terr, y = y_terr, colorscale=cs_name,
                                        name = "Terrain", showscale=False,
                                        contours = {"z" : {"show" : True, "start": 0, "end": np.max(z_terr), "size": 5}}))
    else:
        # add surface plot data
        if show_heatmap:
            plot_data.append(go.Surface(z = z_terr, x = x_terr, y = y_terr, colorscale=cs_name,
                                        surfacecolor=heatmap_data, name = "Terrain", showscale=False))
        else:
            plot_data.append(go.Surface(z = z_terr, x = x_terr, y = y_terr, colorscale=cs_name,
                                        name = "Terrain", showscale=False))

    # add searcher path data
    for i,srch in enumerate(searcher_paths):
        plot_data.append(go.Scatter3d(z = [elv + 1 for elv in srch.smd_history[2]], x = srch.smd_history[0],
                                      y = srch.smd_history[1], mode = "lines+markers",
                                      line=dict(color = "gray", width = 3), name="Searcher {}".format(i+1),
                                      marker=dict(color='gray', size = 2)))
        # add first points
        plot_data.append(go.Scatter3d(z = [srch.smd_history[2][0]], x = [srch.smd_history[0][0]],
                                      y = [srch.smd_history[1][0]], mode = "lines+markers", showlegend=False,
                                      line=dict(color = "gray", width = 3), name="Searcher {}".format(i+1),
                                      marker=dict(color='light gray', size = 12, line=dict(color='black',width=8))))

    # get linear feature info
    [lf_points, inac_points] = linear_features(parameters)
    lf_points = np.vstack([lf_points, z_terr_func.ev(lf_points[0], lf_points[1])+1])
    inac_points = np.vstack([inac_points, z_terr_func.ev(inac_points[0], inac_points[1])+1])

    plot_data.append(go.Scatter3d(z = lf_points[2], x = lf_points[0], y = lf_points[1],
                                  mode = "markers", name="L. Feat.", marker=dict(size = 5)))

    plot_data.append(go.Scatter3d(z = inac_points[2], x = inac_points[0], y = inac_points[1],
                                  mode = "markers", name="Inac. Areas", marker=dict(size = 5)))

    fig = go.Figure(data = plot_data)
    fig.update_layout(title=title, autosize=False, width=2000, height=1200,
                        margin=dict(l=100, r=100, b=65, t=90),
                        legend=dict(yanchor="top", y=0.99, xanchor="left", x=0.01),
                        scene_aspectmode='manual',
                        scene_aspectratio=dict(x=1, y=1, z=0.5))
    fig.show()


if __name__ == "__main__":
    # print("loading data")

    with open("bozoo_save_[37.476914, -80.865864].pkl",'rb') as f:
        loaded_data = pkl.load(f)

    plot_all(parameters = loaded_data['params'],
             mc_object = loaded_data['mc_object'],
             robot_paths = loaded_data['robot_paths'],
             searcher_paths = loaded_data['searcher_paths'],
             title = "Whitethorn, [37.210264, -80.559236]",
             show_heatmap = False,
             show_contours = True,
             cs_name = 'earth')

    with open("bozoo_save_[37.474015, -80.868333].pkl",'rb') as f:
        loaded_data = pkl.load(f)

    plot_all(parameters = loaded_data['params'],
             mc_object = loaded_data['mc_object'],
             robot_paths = loaded_data['robot_paths'],
             searcher_paths = loaded_data['searcher_paths'],
             title = "Mt. Rogers, [36.65711, -81.543333]",
             show_heatmap = False,
             show_contours = True,
             cs_name = 'earth')

    with open("bozoo_save_[37.476261, -80.87046].pkl",'rb') as f:
        loaded_data = pkl.load(f)

    plot_all(parameters = loaded_data['params'],
             mc_object = loaded_data['mc_object'],
             robot_paths = loaded_data['robot_paths'],
             searcher_paths = loaded_data['searcher_paths'],
             title = "Bozoo [37.476914, -80.865864]",
             show_heatmap = False,
             show_contours = True,
             cs_name = 'earth')
