#!/usr/bin/env python
# -*- coding: utf-8 -*-
import argparse
import os

import numpy as np
import plotly.express as px
import plotly.graph_objs as go
import plotly.io as pio
import rasterio
from matplotlib.colors import LightSource, LinearSegmentedColormap


def main(f_path, f_name, farm_name):
    """Main function to visualise the soil moisture content for a given layer,
    year file, and days.

    Parameters
    ----------
    f_path : str
        Path to the file.
    f_name : str
        Name of the file.
    farm_name : str
        Name of the farm.

    Returns
    -------
    None
    """

    # f_path = ""#"/home/sahand/Projects/PIPE-3788 GRDC SoilWaterNow Deployment/work/v3/Arawa 2019-2023/c8/1af25ced023e58c46f4403a155210d/soilwatermodel v3/analysis"
    # f_name = ""#"delta-2022-12-20_2023-01-10-SM2-w_mean-h_mean.tif"
    file = os.path.join(f_path, f_name)
    # farm_name = "Arawa"
    # layer = "SM2"

    with rasterio.open(file) as src:
        band1 = src.read(1)
        print("Band1 has shape", band1.shape)
        height = band1.shape[0]
        width = band1.shape[1]
        cols, rows = np.meshgrid(np.arange(width), np.arange(height))
        xs, ys = rasterio.transform.xy(src.transform, rows, cols)
        lons = np.array(xs)
        lats = np.array(ys)

    bamako = [
        [0.0, "rgb(0, 63, 76)"],  # a scientific colorscale for dem data
        [0.1, "rgb(29, 81, 59)"],
        [0.2, "rgb(55, 98, 43)"],
        [0.3, "rgb(79, 114, 30)"],
        [0.4, "rgb(103, 129, 16)"],
        [0.5, "rgb(136, 142, 2)"],
        [0.6, "rgb(169, 154, 21)"],
        [0.7, "rgb(192, 171, 45)"],
        [0.8, "rgb(214, 188, 74)"],
        [0.9, "rgb(234, 209, 112)"],
        [1.0, "rgb(254, 229, 152)"],
    ]

    # fig= px.imshow(band1, color_continuous_scale=bamako)
    # my_layout= dict(title_text=f"{farm_name}--{f_name.split('.')[0]}", title_x=0.5, width =700, height=500, template='none',
    #                   coloraxis_colorbar=dict(len=0.75, thickness=25))
    # fig.update_layout(**my_layout)

    ve = 1
    ls = LightSource(azdeg=315, altdeg=45)
    hilsh = ls.hillshade(band1, vert_exag=ve)  # , dx=dx, dy=dy)
    fig = px.imshow(hilsh, color_continuous_scale=bamako)

    # fig.update_traces(customdata=elevation, hovertemplate='moisture: %{customdata}<extra></extra>')
    # fig.update_layout(**my_layout)
    # fig.show()#

    cm_data = [
        [0.0011753, 0.25004, 0.3],
        [0.0039003, 0.25157, 0.29861],
        [0.006602, 0.25305, 0.29722],
        [0.0092914, 0.25456, 0.29581],
        [0.012175, 0.25604, 0.2944],
        [0.014801, 0.25755, 0.293],
        [0.01745, 0.25907, 0.29161],
        [0.020096, 0.26057, 0.29018],
        [0.022743, 0.26208, 0.28878],
        [0.025398, 0.26361, 0.28735],
        [0.028064, 0.26511, 0.28592],
        [0.030746, 0.26664, 0.2845],
        [0.033437, 0.26816, 0.28309],
        [0.036369, 0.26971, 0.28165],
        [0.039136, 0.27124, 0.28021],
        [0.041941, 0.27278, 0.27878],
        [0.044597, 0.27432, 0.27734],
        [0.047216, 0.27589, 0.27591],
        [0.049816, 0.27743, 0.27443],
        [0.052305, 0.279, 0.27298],
        [0.054846, 0.28053, 0.27151],
        [0.057331, 0.28213, 0.27007],
        [0.059903, 0.28369, 0.2686],
        [0.062324, 0.28526, 0.2671],
        [0.064737, 0.28684, 0.26564],
        [0.067177, 0.28844, 0.26416],
        [0.069647, 0.29003, 0.26267],
        [0.072007, 0.29165, 0.26118],
        [0.074435, 0.29324, 0.25969],
        [0.076803, 0.29484, 0.25817],
        [0.079232, 0.29647, 0.25667],
        [0.081738, 0.29808, 0.25518],
        [0.084163, 0.29972, 0.25365],
        [0.086556, 0.30134, 0.25213],
        [0.089008, 0.30298, 0.25059],
        [0.091487, 0.30463, 0.24907],
        [0.093906, 0.3063, 0.24754],
        [0.096327, 0.30793, 0.24598],
        [0.098845, 0.30961, 0.24443],
        [0.10132, 0.31127, 0.24288],
        [0.10384, 0.31294, 0.24134],
        [0.10637, 0.31461, 0.23978],
        [0.10884, 0.31632, 0.23821],
        [0.1114, 0.318, 0.23666],
        [0.11391, 0.3197, 0.23509],
        [0.11648, 0.32139, 0.23347],
        [0.11904, 0.32311, 0.23191],
        [0.12156, 0.32481, 0.2303],
        [0.12416, 0.32653, 0.22873],
        [0.12676, 0.32826, 0.22712],
        [0.12938, 0.33, 0.22551],
        [0.132, 0.33173, 0.22389],
        [0.13465, 0.33348, 0.22228],
        [0.13728, 0.33524, 0.22068],
        [0.13988, 0.33699, 0.21904],
        [0.14258, 0.33874, 0.21741],
        [0.14528, 0.34051, 0.21578],
        [0.14793, 0.34229, 0.21413],
        [0.15062, 0.34406, 0.21248],
        [0.15331, 0.34584, 0.21083],
        [0.15607, 0.34765, 0.20917],
        [0.15879, 0.34944, 0.20751],
        [0.16154, 0.35125, 0.20585],
        [0.16429, 0.35306, 0.20419],
        [0.16705, 0.35489, 0.20248],
        [0.16983, 0.35671, 0.20078],
        [0.17259, 0.35853, 0.1991],
        [0.1754, 0.36036, 0.19742],
        [0.17824, 0.36222, 0.19572],
        [0.18102, 0.36406, 0.19403],
        [0.18387, 0.36593, 0.1923],
        [0.18672, 0.36779, 0.19058],
        [0.18959, 0.36966, 0.18887],
        [0.19246, 0.37154, 0.18714],
        [0.19535, 0.37344, 0.18542],
        [0.19825, 0.37533, 0.18364],
        [0.20116, 0.37724, 0.18189],
        [0.20413, 0.37915, 0.18014],
        [0.20706, 0.38108, 0.17841],
        [0.21001, 0.383, 0.1766],
        [0.21297, 0.38493, 0.17486],
        [0.21599, 0.38688, 0.17306],
        [0.21898, 0.38884, 0.17127],
        [0.22199, 0.39081, 0.1695],
        [0.22501, 0.39277, 0.16768],
        [0.22809, 0.39476, 0.16585],
        [0.23115, 0.39674, 0.16407],
        [0.23423, 0.39875, 0.16224],
        [0.23735, 0.40078, 0.16037],
        [0.24043, 0.40279, 0.15857],
        [0.24357, 0.40482, 0.15671],
        [0.24673, 0.40687, 0.15484],
        [0.24989, 0.40893, 0.15296],
        [0.2531, 0.411, 0.15109],
        [0.25629, 0.41308, 0.14922],
        [0.25954, 0.41517, 0.14733],
        [0.26278, 0.41728, 0.14546],
        [0.26606, 0.41939, 0.14352],
        [0.26935, 0.42151, 0.14161],
        [0.27265, 0.42366, 0.13963],
        [0.276, 0.42581, 0.13775],
        [0.27935, 0.42799, 0.13575],
        [0.28274, 0.43017, 0.13382],
        [0.28611, 0.43238, 0.13185],
        [0.28956, 0.43459, 0.12989],
        [0.293, 0.43682, 0.12788],
        [0.29648, 0.43906, 0.12588],
        [0.29998, 0.44132, 0.12382],
        [0.3035, 0.44361, 0.12179],
        [0.30706, 0.44589, 0.11977],
        [0.31066, 0.44821, 0.11773],
        [0.31426, 0.45054, 0.11568],
        [0.3179, 0.45289, 0.11362],
        [0.32156, 0.45525, 0.11153],
        [0.32526, 0.45763, 0.10944],
        [0.32901, 0.46003, 0.10731],
        [0.33277, 0.46245, 0.10513],
        [0.33654, 0.4649, 0.10304],
        [0.34037, 0.46736, 0.10085],
        [0.34422, 0.46984, 0.098693],
        [0.34812, 0.47233, 0.096453],
        [0.35204, 0.47486, 0.094325],
        [0.35601, 0.47738, 0.092102],
        [0.36, 0.47994, 0.089849],
        [0.36403, 0.48251, 0.087617],
        [0.36811, 0.4851, 0.085276],
        [0.37221, 0.48772, 0.082993],
        [0.37638, 0.49034, 0.080686],
        [0.38057, 0.49296, 0.078304],
        [0.3848, 0.49562, 0.075955],
        [0.38909, 0.49829, 0.073656],
        [0.39341, 0.50095, 0.071224],
        [0.39779, 0.50362, 0.068802],
        [0.4022, 0.5063, 0.066264],
        [0.40668, 0.50898, 0.063774],
        [0.4112, 0.51164, 0.061172],
        [0.41578, 0.51431, 0.058651],
        [0.4204, 0.51695, 0.05607],
        [0.42509, 0.51956, 0.053359],
        [0.42981, 0.52214, 0.050712],
        [0.4346, 0.5247, 0.047969],
        [0.43943, 0.52722, 0.045272],
        [0.44432, 0.52967, 0.042483],
        [0.44924, 0.53207, 0.039697],
        [0.45421, 0.5344, 0.036906],
        [0.45922, 0.53666, 0.034001],
        [0.46428, 0.53884, 0.031435],
        [0.46935, 0.54092, 0.028957],
        [0.47444, 0.54291, 0.026597],
        [0.47956, 0.5448, 0.024363],
        [0.48468, 0.54659, 0.022265],
        [0.48982, 0.54827, 0.020312],
        [0.49497, 0.54985, 0.018512],
        [0.5001, 0.55132, 0.016876],
        [0.50524, 0.55269, 0.015412],
        [0.51038, 0.55397, 0.014132],
        [0.5155, 0.55517, 0.013033],
        [0.52063, 0.55628, 0.01218],
        [0.52575, 0.55735, 0.011416],
        [0.53088, 0.55835, 0.010839],
        [0.53601, 0.55933, 0.010585],
        [0.54115, 0.5603, 0.010612],
        [0.54633, 0.56126, 0.010945],
        [0.55153, 0.56225, 0.011649],
        [0.55677, 0.56328, 0.012516],
        [0.56207, 0.56438, 0.01365],
        [0.56742, 0.56555, 0.015118],
        [0.57285, 0.56684, 0.016927],
        [0.57835, 0.56824, 0.019098],
        [0.58394, 0.56978, 0.021653],
        [0.58963, 0.57149, 0.024617],
        [0.59541, 0.57336, 0.028014],
        [0.60129, 0.5754, 0.03187],
        [0.60728, 0.57763, 0.036408],
        [0.61335, 0.58006, 0.041287],
        [0.61954, 0.58268, 0.046299],
        [0.62581, 0.58551, 0.051565],
        [0.63218, 0.58853, 0.057023],
        [0.63862, 0.59176, 0.062625],
        [0.64514, 0.59517, 0.068288],
        [0.65173, 0.59876, 0.074092],
        [0.65837, 0.60254, 0.079969],
        [0.66506, 0.60647, 0.085984],
        [0.6718, 0.61057, 0.0922],
        [0.67856, 0.6148, 0.098456],
        [0.68534, 0.61918, 0.10476],
        [0.69213, 0.62366, 0.11129],
        [0.69891, 0.62825, 0.11783],
        [0.70569, 0.63293, 0.12447],
        [0.71245, 0.63769, 0.13125],
        [0.71918, 0.64252, 0.1381],
        [0.72588, 0.6474, 0.14503],
        [0.73253, 0.65232, 0.15201],
        [0.73912, 0.65728, 0.15912],
        [0.74564, 0.66226, 0.16629],
        [0.7521, 0.66724, 0.17352],
        [0.75847, 0.67223, 0.18082],
        [0.76475, 0.67721, 0.18823],
        [0.77094, 0.68215, 0.19562],
        [0.77702, 0.68709, 0.20308],
        [0.78299, 0.69199, 0.21059],
        [0.78886, 0.69684, 0.21813],
        [0.79459, 0.70164, 0.2257],
        [0.80021, 0.7064, 0.23325],
        [0.8057, 0.7111, 0.24084],
        [0.81106, 0.71574, 0.24845],
        [0.81631, 0.72031, 0.25601],
        [0.82142, 0.72482, 0.26361],
        [0.8264, 0.72926, 0.27115],
        [0.83127, 0.73364, 0.2787],
        [0.83602, 0.73795, 0.28619],
        [0.84066, 0.74221, 0.29369],
        [0.84519, 0.74639, 0.30114],
        [0.84961, 0.75052, 0.30859],
        [0.85394, 0.75459, 0.31598],
        [0.85818, 0.7586, 0.32333],
        [0.86234, 0.76255, 0.33065],
        [0.86641, 0.76646, 0.33794],
        [0.87042, 0.77032, 0.3452],
        [0.87436, 0.77414, 0.35241],
        [0.87824, 0.77792, 0.3596],
        [0.88207, 0.78167, 0.36676],
        [0.88585, 0.78537, 0.37388],
        [0.88958, 0.78905, 0.38098],
        [0.89328, 0.7927, 0.38804],
        [0.89694, 0.79632, 0.39508],
        [0.90057, 0.79992, 0.40209],
        [0.90417, 0.8035, 0.40909],
        [0.90775, 0.80706, 0.41606],
        [0.9113, 0.8106, 0.423],
        [0.91484, 0.81412, 0.42995],
        [0.91835, 0.81764, 0.43686],
        [0.92186, 0.82114, 0.44376],
        [0.92534, 0.82462, 0.45064],
        [0.92882, 0.8281, 0.45751],
        [0.93228, 0.83157, 0.46438],
        [0.93573, 0.83502, 0.47121],
        [0.93918, 0.83847, 0.47804],
        [0.94261, 0.84192, 0.48486],
        [0.94603, 0.84535, 0.49169],
        [0.94945, 0.84878, 0.49849],
        [0.95286, 0.85221, 0.50528],
        [0.95627, 0.85563, 0.51207],
        [0.95966, 0.85904, 0.51886],
        [0.96305, 0.86246, 0.52563],
        [0.96643, 0.86586, 0.53241],
        [0.96981, 0.86927, 0.53917],
        [0.97318, 0.87268, 0.54594],
        [0.97654, 0.87608, 0.55269],
        [0.9799, 0.87948, 0.55945],
        [0.98325, 0.88288, 0.5662],
        [0.98659, 0.88628, 0.57296],
        [0.98993, 0.88968, 0.57971],
        [0.99326, 0.89308, 0.58645],
        [0.99658, 0.89648, 0.5932],
        [0.9999, 0.89988, 0.59995],
    ]

    bamako_map = LinearSegmentedColormap.from_list("bamako", cm_data)
    rgb = ls.shade(
        band1, cmap=bamako_map, blend_mode="overlay", vert_exag=ve
    )  # , dx=dx, dy=dy)
    img = np.array((255 * rgb[:, :, :3]), int)

    fig = go.Figure(
        go.Image(
            z=img,
            colormodel="rgb",
            customdata=np.stack((band1, lats, lons), axis=-1),
            hovertemplate="<b>SM</b>: %{customdata[0]}<br>"
            + "<b>Lat</b>: %{customdata[1]}<br>"
            + "<b>Lon</b>: %{customdata[2]}<br>"
            + "<extra></extra>",
        )
    )
    fig.update_layout(
        title_text=f"{farm_name}--{f_name.split('.')[0]}",
        title_x=0.5,
        width=700,
        height=500,
        template="none",
        yaxis_autorange="reversed",
    )
    # fig.show()
    pio.write_html(fig, file=f"{file.replace('.tif','')}plot.html", auto_open=True)


if __name__ == "__main__":
    # Load Configs
    parser = argparse.ArgumentParser(
        description="Perform mass conservation analysis on the soilwatermodel logs. This should be called after running the model module with `-m True` argument.",
        formatter_class=argparse.ArgumentDefaultsHelpFormatter,
    )
    parser.add_argument(
        "-f",
        "--file",
        help="file name",
    )
    parser.add_argument(
        "-d",
        "--dir",
        help="directory name",
    )
    parser.add_argument(
        "-n",
        "--name",
        help="farm name",
    )

    args = parser.parse_args()
    f_name = args.file
    d_name = args.dir
    farm_name = args.name

    main(f_name, d_name, farm_name)