# ---
# jupyter:
#   jupytext:
#     formats: ipynb,py:percent
#     text_representation:
#       extension: .py
#       format_name: percent
#       format_version: '1.3'
#       jupytext_version: 1.16.2
#   kernelspec:
#     display_name: temps
#     language: python
#     name: temps
# ---

# %% [markdown]
# # FIGURE METRICS

# %% [markdown]
# ## METRICS FOR THE DIFFERENT METHODS ON THE WIDE FIELD SAMPLE

# %% [markdown]
# ### LOAD PYTHON MODULES

# %%
# %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"


# %%
import temps

# %%
from temps.archive import Archive
from temps.utils import nmad
from temps.temps_arch import EncoderPhotometry, MeasureZ
from temps.temps import TempsModule
from temps.plots import plot_photoz


# %%
eval_methods = True

# %% [markdown]
# ### 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"

# %%
path_file = parent_dir / filename_valid  # Creating the path to the file
hdu_list = fits.open(path_file)
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)

# %% [markdown]
# ### EVALUATE USING TRAINED MODELS

# %%
if eval_methods:

    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, return_flag=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


# %%
dfs["z"]["zwerr"] = (dfs["z"].z - dfs["z"].ztarget) / (1 + dfs["z"].ztarget)
dfs["L15"]["zwerr"] = (dfs["L15"].z - dfs["L15"].ztarget) / (1 + dfs["L15"].ztarget)
dfs["DA"]["zwerr"] = (dfs["DA"].z - dfs["DA"].ztarget) / (1 + dfs["DA"].ztarget)

# %% [markdown]
# ### LOAD CATALOGUES FROM PREVIOUS TRAINING

# %%
if not eval_methods:
    dfs = {}
    dfs["z"] = pd.read_csv(parent_dir / "predictions_specztraining.csv", header=0)
    dfs["L15"] = pd.read_csv(parent_dir / "predictions_speczL15training.csv", header=0)
    dfs["DA"] = pd.read_csv(parent_dir / "predictions_speczDAtraining.csv", header=0)


# %% [markdown]
# ### MAKE PLOT

# %%
df_list = [dfs["z"], dfs["L15"], dfs["DA"]]

# %%
plot_photoz(
    df_list,
    nbins=8,
    xvariable="VISmag",
    metric="nmad",
    type_bin="bin",
    label_list=["zs", "zs+L15", r"TEMPS"],
    save=False,
    samp="L15",
)

# %%