# ---
# jupyter:
#   jupytext:
#     text_representation:
#       extension: .py
#       format_name: light
#       format_version: '1.5'
#       jupytext_version: 1.16.2
#   kernelspec:
#     display_name: temps
#     language: python
#     name: temps
# ---

# # FIGURE COLOURSPACE IN THE PAPER

# %load_ext autoreload
# %autoreload 2

import pandas as pd
import numpy as np
import os
from astropy.io import fits
from astropy.table import Table
import torch
from pathlib import Path

#matplotlib settings
from matplotlib import rcParams
import matplotlib.pyplot as plt
rcParams["mathtext.fontset"] = "stix"
rcParams["font.family"] = "STIXGeneral"

from temps.archive import Archive 
from temps.utils import nmad
from temps.temps_arch import EncoderPhotometry, MeasureZ
from temps.temps import TempsModule


def estimate_som_map(df, plot_arg='z', nx=40, ny=40):
    """
    Estimate a Self-Organizing Map (SOM) visualization from a DataFrame.

    Parameters:
    - df (pandas.DataFrame): Input DataFrame containing data for SOM estimation.
    - plot_arg (str, optional): Column name to be used for plotting. Default is 'z'.
    - nx (int, optional): Number of cells along the X-axis. Default is 40.
    - ny (int, optional): Number of cells along the Y-axis. Default is 40.

    Returns:
    - som_data (numpy.ndarray): Estimated SOM visualization data.
    """
    x_cells = np.arange(0, nx)
    y_cells = np.arange(0, ny)
    index_cell = np.arange(nx * ny)
    cells = np.array(np.meshgrid(x_cells, y_cells)).T.reshape(-1, 2)
    cells = pd.DataFrame(np.c_[cells[:, 0], cells[:, 1], index_cell], columns=['x_cell', 'y_cell', 'cell'])

    if plot_arg == 'count':
        som_vis = df.groupby('cell')['z'].count().reset_index().rename(columns={f'z': 'plot_som'})
    else:
        som_vis = df.groupby('cell')[f'{plot_arg}'].mean().reset_index().rename(columns={f'{plot_arg}': 'plot_som'})

    som_data = som_vis.merge(cells, on='cell')
    som_data = som_data.pivot(index='x_cell', columns='y_cell', values='plot_som')

    return som_data



def plot_som_map(som_data, plot_arg = 'z', vmin=0, vmax=1):
    """
    Plot the Self-Organizing Map (SOM) data.

    Parameters:
    - som_data (numpy.ndarray): The SOM data to be visualized.
    - plot_arg (str, optional): The column name to be plotted. Default is 'z'.
    - vmin (float, optional): Minimum value for color scaling. Default is 0.
    - vmax (float, optional): Maximum value for color scaling. Default is 1.

    Returns:
    None
    """
    plt.imshow(som_data, vmin=vmin, vmax=vmax, cmap='viridis')  # Choose an appropriate colormap
    plt.colorbar(label=f'{plot_arg}')  # Add a colorbar with a label
    plt.xlabel(r'$x$ [pixel]', fontsize=14)  # Add an appropriate X-axis label
    plt.ylabel(r'$y$ [pixel]', fontsize=14)  # Add an appropriate Y-axis label
    plt.show()



# ### LOAD DATA

#define here the directory containing the photometric catalogues
parent_dir = Path('/data/astro/scratch/lcabayol/insight/data/Euclid_EXT_MER_PHZ_DC2_v1.5')
modules_dir = Path('../data/models/')
filename_calib = 'euclid_cosmos_DC2_S1_v2.1_calib_clean.fits'
filename_valid = 'euclid_cosmos_DC2_S1_v2.1_valid_matched.fits'

# +
filename_valid='euclid_cosmos_DC2_S1_v2.1_valid_matched.fits'

hdu_list = fits.open(parent_dir/filename_valid)
cat = Table(hdu_list[1].data).to_pandas()
cat = cat[cat['FLAG_PHOT']==0]
cat = cat[cat['mu_class_L07']==1]
cat = cat[(cat['z_spec_S15'] > 0) | (cat['photo_z_L15'] > 0)]
cat = cat[cat['MAG_VIS']<25]

# -

ztarget = [cat['z_spec_S15'].values[ii] if cat['z_spec_S15'].values[ii]> 0 else cat['photo_z_L15'].values[ii] for ii in range(len(cat))]
specz_or_photo = [0 if cat['z_spec_S15'].values[ii]> 0 else 1 for ii in range(len(cat))]
ID = cat['ID']
VISmag = cat['MAG_VIS']
zsflag = cat['reliable_S15']

photoz_archive = Archive(path_calib = parent_dir/filename_calib, 
                         path_valid = parent_dir/filename_valid,
                         only_zspec=False)
f = photoz_archive._extract_fluxes(catalogue= cat)
col = photoz_archive._to_colors(f)

device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')

# +
dfs = {}

for il, lab in enumerate(['z','L15','DA']):

    nn_features = EncoderPhotometry()
    nn_features.load_state_dict(torch.load(modules_dir / f'modelF_{lab}.pt',map_location=torch.device('cpu')))
    nn_z = MeasureZ(num_gauss=6)
    nn_z.load_state_dict(torch.load(modules_dir / f'modelZ_{lab}.pt',map_location=torch.device('cpu')))
    
    temps_module = TempsModule(nn_features, nn_z)
    
    z, pz, odds = temps_module.get_pz(input_data=torch.Tensor(col), 
                                return_pz=True)
    # Create a DataFrame with the desired columns
    df = pd.DataFrame(np.c_[ID, VISmag,z,odds, ztarget,zsflag, specz_or_photo], 
                      columns=['ID','VISmag','z', 'odds','ztarget','zsflag','S15_L15_flag'])

    # Calculate additional columns or operations if needed
    df['zwerr'] = (df.z - df.ztarget) / (1 + df.ztarget)

    # Drop any rows with NaN values
    df = df.dropna()

    # Assign the DataFrame to a key in the dictionary
    dfs[lab] = df

# -

# ### LOAD TRAINED MODELS AND EVALUATE PDFs AND REDSHIFT

#define here the directory containing the photometric catalogues
parent_dir = Path('/data/astro/scratch/lcabayol/insight/data/Euclid_EXT_MER_PHZ_DC2_v1.5')
modules_dir = Path('../data/models/')

df_z = dfs['z']
df_z_DA = dfs['DA']

# ##### LOAD TRAIN SOM ON TRAINING DATA

df_som = pd.read_csv(parent_dir/'som_dataframe.csv', header = 0, sep =',')
df_z = df_z.merge(df_som, on = 'ID')
df_z_DA = df_z_DA.merge(df_som, on = 'ID')

# ##### APPLY CUTS FOR DIFFERENT SAMPLES

df_zspec = df_z[(df_z.S15_L15_flag==0) & (df_z.zsflag==1)]
df_l15 = df_z[(df_z.ztarget>0)]
df_l15_DA = df_z_DA[(df_z_DA.ztarget>0)]

df_l15_euclid = df_z[(df_z.VISmag <24.5) & (df_z.z > 0.2) & (df_z.z < 2.6)]
df_l15_euclid_cut= df_l15_euclid[df_l15_euclid.odds>df_l15_euclid['odds'].quantile(0.2)]

df_l15_euclid_da = df_z_DA[(df_z_DA.VISmag <24.5) & (df_z_DA.z > 0.2) & (df_z_DA.z < 2.6)]
df_l15_euclid_cut_da= df_l15_euclid_da[df_l15_euclid_da.odds>df_l15_euclid['odds'].quantile(0.2)]

# ## MAKE SOM PLOT

from mpl_toolkits.axes_grid1 import make_axes_locatable

# +
fig, axs = plt.subplots(6, 4, figsize=(13, 15), sharex=True, sharey=True, gridspec_kw={'hspace': 0.05, 'wspace': 0.06})

# Plot in the top row (axs[0, i])
#top row, spectroscopic sample
columns = ['ztarget','z','zwerr','count']
titles = [r'$z_{true}$ (A)',r'$z$ (B)',r'$z_{\rm error}$ (C)','Counts']
limits = [[0,4],[0,4],[-0.5,0.5],[0,50]]
for ii in range(4):
    som_data = estimate_som_map(df_zspec, plot_arg=columns[ii], nx=40, ny=40)
    im = axs[0,ii].imshow(som_data, vmin=limits[ii][0], vmax=limits[ii][1], cmap='viridis')  # Choose an appropriate colormap
    axs[0, ii].set_title(f'{titles[ii]}', fontsize=18)
    
    if ii==0:
        axs[0, 0].set_ylabel(r'$y$', fontsize=14)
    elif ii==1:
        cbar_ax = fig.add_axes([0.49, 0.11, 0.01, 0.77])
        fig.colorbar(im, cax=cbar_ax)
    elif ii==2:
        cbar_ax = fig.add_axes([0.685, 0.11, 0.01, 0.77])
        fig.colorbar(im, cax=cbar_ax)
    elif ii==3:
        cbar_ax = fig.add_axes([0.885, 0.11, 0.01, 0.77])
        fig.colorbar(im, cax=cbar_ax)

for jj in range(4):
    som_data = estimate_som_map(df_l15, plot_arg=columns[jj], nx=40, ny=40)
    im = axs[1,jj].imshow(som_data, vmin=limits[jj][0], vmax=limits[jj][1], cmap='viridis')  # Choose an appropriate colormap
    #axs[1, jj].set_title(f'{titles[jj]}', fontsize=14)
    #axs[1, jj].set_xlabel(r'$x$', fontsize=14)
    
    
for kk in range(4):
    som_data = estimate_som_map(df_l15_DA, plot_arg=columns[kk], nx=40, ny=40)
    im = axs[2,kk].imshow(som_data, vmin=limits[kk][0], vmax=limits[kk][1], cmap='viridis')  # Choose an appropriate colormap
    #axs[2, kk].set_title(f'{titles[kk]}', fontsize=14)
    #axs[2, kk].set_xlabel(r'$x$', fontsize=14)
    
for rr in range(4):
    som_data = estimate_som_map(df_l15_euclid_da, plot_arg=columns[rr], nx=40, ny=40)
    im = axs[3,rr].imshow(som_data, vmin=limits[rr][0], vmax=limits[rr][1], cmap='viridis')  # Choose an appropriate colormap
    #axs[3, rr].set_title(f'{titles[rr]}', fontsize=14)
    #axs[3, rr].set_xlabel(r'$x$', fontsize=14)
    
for ll in range(4):
    som_data = estimate_som_map(df_l15_euclid_cut, plot_arg=columns[ll], nx=40, ny=40)
    im = axs[4,ll].imshow(som_data, vmin=limits[ll][0], vmax=limits[ll][1], cmap='viridis')  # Choose an appropriate colormap
    #axs[4, ll].set_title(f'{titles[ll]}', fontsize=14)
    axs[4, ll].set_xlabel(r'$x$', fontsize=14)
    
for ll in range(4):
    som_data = estimate_som_map(df_l15_euclid_cut_da, plot_arg=columns[ll], nx=40, ny=40)
    im = axs[5,ll].imshow(som_data, vmin=limits[ll][0], vmax=limits[ll][1], cmap='viridis')  # Choose an appropriate colormap
    #axs[4, ll].set_title(f'{titles[ll]}', fontsize=14)
    axs[5, ll].set_xlabel(r'$x$', fontsize=14)

    
axs[0, 0].set_ylabel(r'$y$', fontsize=14)
axs[1, 0].set_ylabel(r'$y$', fontsize=14)
axs[2, 0].set_ylabel(r'$y$', fontsize=14)
axs[3, 0].set_ylabel(r'$y$', fontsize=14)
axs[4, 0].set_ylabel(r'$y$', fontsize=14)
axs[5, 0].set_ylabel(r'$y$', fontsize=14)


fig.text(0.09, 0.815, r'$z_{\rm s}$ samp. (1)', va='center', rotation='vertical', fontsize=16)
fig.text(0.09, 0.69, r'L15 samp. (2)', va='center', rotation='vertical', fontsize=16)
fig.text(0.09, 0.56, r'L15 samp. + DA (3)', va='center', rotation='vertical', fontsize=14)
fig.text(0.09, 0.44, r'$Euclid$ samp. + DA (4)', va='center', rotation='vertical', fontsize=14)
fig.text(0.09, 0.3, r'$Euclid$ samp. + QC (5)', va='center', rotation='vertical', fontsize=14)

fig.text(0.09, 0.17, r'(5) + DA ', va='center', rotation='vertical', fontsize=13)


plt.savefig('SOM_colourspace.pdf', format='pdf', bbox_inches='tight', dpi=300)

# -