# import os
import rasterio
# import earthpy as et
import earthpy.spatial as es
# import earthpy.plot as ep
# import matplotlib as mpl
# from matplotlib import pyplot as plt
# from matplotlib.colors import LinearSegmentedColormap
# import datetime
import numpy as np
# import pandas as pd
import sklearn.decomposition

# import process_sb_tiff
import disco_tif.geotiff_plotting

#################################################

def build_hs_az_al(start_al, al_inc, num_al_angles, start_az, num_az_angles):
    """description
Input Parameters:
    - start_al:
    
    - al_inc:
    
    - num_al_angles:
    
    - start_az:
    
    - num_az_angles
    """
    azimuths = np.linspace(start_az, start_az+360, num_az_angles+1)  # degrees 0=360, number of divisions + 1 because start=end
    filt = azimuths >= 360
    while filt.sum() > 0:
        azimuths[filt] = azimuths[filt]-360
        filt = azimuths >= 360
    azimuths = azimuths[:-1].astype(int).tolist()
    azimuths = sorted(azimuths)
    print(f"azimuths = {azimuths}")
    
    altitudes = np.linspace(start_al, (al_inc*num_al_angles)+start_al, num_al_angles+1)
    altitudes = altitudes[:-1].astype(int).tolist()
    print(f"altitudes = {altitudes}")
    return azimuths, altitudes

def write_raster_dict_data_to_geotiff(single_band_tiff_path, origprofile, raster_data_dict, len_hs=None):
    """description
Input Parameters:
    - single_band_tiff_path:
    
    - origprofile:
    
    - raster_data_dict:
    
    - len_hs=None
    """
    for key, value in raster_data_dict.items():
        if 'component' in key:
            assert len_hs is not None, "'len_hs' cannot be none if passing in a pca_dictionary_object"
            new_tiff_path = f"{single_band_tiff_path.split('.tif')[0]}_hillshade_pca{len_hs}-{key}.tif"
        else:
            new_tiff_path = f"{single_band_tiff_path.split('.tif')[0]}_hillshade_{key}.tif"
        
        newprofile = origprofile.copy()
        newprofile.update(dtype=str((value[0, 0]).dtype), nodata=np.nan)
           
        with rasterio.open(new_tiff_path, 'w', **newprofile) as dst:
            dst.write(arr=value, indexes=1, masked=True)
    
        print(f"New single-channel geotiff generated successfully: '{new_tiff_path}'")
        

def MakeHillShadePCA(hillshades, plot_figures=False, raster_data=None, cmap='terrain', n_components=3):
    """description
Input Parameters:
    - hillshades:
    
    - plot_figures=False:
    
    - raster_data=None:
    
    - cmap='terrain':
    
    - n_components=3
    """
    if plot_figures:
        assert raster_data is not None, "raster data must be supplied if plot_figures is true"
        
    NoNanIndicies = np.argwhere(~np.isnan(hillshades[list(hillshades.keys())[0]].flatten())).flatten()

    n_samples = len(NoNanIndicies)
    n_features = len(hillshades.keys())
    assert n_features >= 4, "There must be at least 4 hillshades to produce a PCA"
    
    flat_hillshades = np.ones([n_samples, n_features]) * np.nan
    
    feature_ind = -1
    for key, value in hillshades.items():
        feature_ind += 1
        value = value.flatten()
        value = value[NoNanIndicies]
        flat_hillshades[:, feature_ind] = value
    
    assert np.isnan(flat_hillshades).sum() == 0, "There can be no NaN's when performing a pca"
    
    pca = sklearn.decomposition.PCA(n_components=n_components)
    
    pcaout = pca.fit(flat_hillshades).transform(flat_hillshades)

    data = hillshades[list(hillshades.keys())[0]]
    
    nrow = data.shape[0]
    ncol = data.shape[1]
    nlay = len(pcaout[0])
    
    dumarray = np.ones([nrow * ncol])*np.nan
    
    pcaComponents = {}
    for ilay in range(0, len(pcaout[0])):
        tdat = dumarray.copy()
        nowdat = pcaout[:, ilay]
        tdat[NoNanIndicies] = nowdat
        pcaComponents[f'component_{ilay+1}'] = tdat.reshape([nrow, ncol], order='C')
    
    if plot_figures:
        # plot the pca outputs
        disco_tif.geotiff_plotting.plot_greyband_only(raster_data_dict=pcaComponents,
                                                      nrows=1,
                                                      ncols=n_components)
        disco_tif.geotiff_plotting.plot_color_raster_with_greyscale_overlay(raster_data=raster_data,
                                                                            cmap=cmap,
                                                                            raster_data_dict=pcaComponents,
                                                                            nrows=1,
                                                                            ncols=n_components)
    
    # pcaComponents['component_1_0_1'] = (           pcaComponents['component_1'] - np.nanmin(pcaComponents['component_1'])) / \
    #                                    (np.nanmax(pcaComponents['component_1']) - np.nanmin(pcaComponents['component_1']))
    
    # pcaComponents['component_2_0_1'] = (           pcaComponents['component_2'] - np.nanmin(pcaComponents['component_2'])) / \
    #                                    (np.nanmax(pcaComponents['component_2']) - np.nanmin(pcaComponents['component_2']))
    
    # pcaComponents['component_3_0_1'] = (           pcaComponents['component_3'] - np.nanmin(pcaComponents['component_3'])) / \
    #                                    (np.nanmax(pcaComponents['component_3']) - np.nanmin(pcaComponents['component_3']))
                                                                                                                   
    return pcaComponents
    

def build_hillshade(single_band_tiff_path, data_min_max,  hs_azimuths, hs_altitudes, cmap='terrain', process_pca=False, plot_figures=False):
    """description
Input Parameters:
    - single_band_tiff_path:
    
    - data_min_max:
    
    - hs_azimuths:
    
    - hs_altitudes:
    
    - cmap='terrain':
    
    - process_pca=False:
    
    - plot_figures=False
    """
    # read geotiff and minimally process for the colormap function
    with rasterio.open(single_band_tiff_path, 'r') as src:
        data = src.read(1)  # Read the first band
        no_data_value = src.nodata  # Get the no-data value from the GeoTIFF
        epsg_code = src.crs.to_epsg() if src.crs else None
        origprofile = src.profile
        width = src.width 
        height = src.height
        srcmask = src.read_masks(1)
    
    # Clip data values to the specified range
    clipped_data = np.clip(data, data_min_max[0], data_min_max[1])
    
    nan_data = data.copy().astype(float)
    nan_clipped_data = clipped_data.copy().astype(float)
    
    if no_data_value is not None:
        nan_data[data == no_data_value] = np.nan
        nan_clipped_data[data == no_data_value] = np.nan

    #
    if plot_figures:
        # Plot the data
        disco_tif.geotiff_plotting.plot_singleband_raster(raster_data=nan_clipped_data,
                                                          cmap=cmap,
                                                          title=f"{single_band_tiff_path} Without Hillshade",
                                                          ax=None)
       
    # calculate the hillshade for  all combination fo azimuths and altitudes
    num_hillshades = len(hs_azimuths) * len(hs_altitudes)
    # it it expensive to generate all of these hillshades and do the pca on them => we limit the number to something reasonable # 8 azimuths at 3 angles will produce 24 hillshades. ### what is reasonable here?
    max_num_hs = 24
    assert num_hillshades <= max_num_hs, f"Only {max_num_hs} azimuth-altitude combinations can be used to calculate a Hillshade PCA but {num_hillshades} were provided"

    hillshades = {}
    for az_ind, my_azimuth in enumerate(hs_azimuths):
        for al_ind, my_altitude in enumerate(hs_altitudes):
            # Create and plot the hillshade with earthpy
            txt_my_azimuth = f"00{my_azimuth}"[-3:]
            txt_my_altitude = f"0{my_altitude}"[-2:]
            hskey = f'az{txt_my_azimuth}_al{txt_my_altitude}'
            hillshades[hskey] = es.hillshade(nan_data, azimuth=my_azimuth, altitude=my_altitude)

    if plot_figures:
        # plot the hillshades 
        disco_tif.geotiff_plotting.plot_greyband_only(raster_data_dict=hillshades,
                                                      nrows=max([len(hs_azimuths), len(hs_altitudes)]),
                                                      ncols=min([len(hs_azimuths), len(hs_altitudes)]))
        # plot geotiff overlain with the hillshades
        disco_tif.geotiff_plotting.plot_color_raster_with_greyscale_overlay(raster_data=nan_clipped_data,
                                                                            raster_data_dict=hillshades,
                                                                            nrows=max([len(hs_azimuths), len(hs_altitudes)]),
                                                                            ncols=min([len(hs_azimuths), len(hs_altitudes)]),
                                                                            cmap=cmap)

    write_raster_dict_data_to_geotiff(single_band_tiff_path, origprofile, hillshades)

    if not process_pca:
        return hillshades
    elif process_pca:
        assert num_hillshades >= 4, f'Need at least 4 azimuth-altitude combinations, only {num_hillshades} were provided'
        pcaComponents = MakeHillShadePCA(hillshades, plot_figures, raster_data=nan_clipped_data, cmap=cmap)
        write_raster_dict_data_to_geotiff(single_band_tiff_path, origprofile, pcaComponents, len_hs=len(hillshades.keys()))
        return hillshades, pcaComponents
