import torch
import warnings
import numpy as np
import pandas as pd
from pathlib import Path
from ctapipe.io import TableLoader
from scripts.make_graphs import make_graphs




def to_log_transform(df):
    """
    Determine which features in the DataFrame should be log-transformed.
    Args:
        df (pandas.DataFrame): The input DataFrame containing the features.
    Returns:
        list: A list of feature names that should be log-transformed.
    """

    FEATURES_TO_LOG_TRANSFORM = ["true_energy", "hillas_intensity"]
    if "RandomForestRegressor_tel_energy" in df.columns:
        FEATURES_TO_LOG_TRANSFORM.append("RandomForestRegressor_tel_energy")
    return FEATURES_TO_LOG_TRANSFORM


def log_transform(df):
    """
    Applies a log10 transformation to specified features in the DataFrame.
    Parameters:
    df (pandas.DataFrame): The input DataFrame containing the features to be transformed.
    Returns:
    pandas.DataFrame: The DataFrame with the specified features log10 transformed.
    """

    # apply on FEATURES_TO_LOG_TRANSFORM
    FEATURES_TO_LOG_TRANSFORM = to_log_transform(df)
    df[FEATURES_TO_LOG_TRANSFORM] = np.log10(df[FEATURES_TO_LOG_TRANSFORM])

    return df


def inverse_log_transform(df):

    # apply on FEATURES_TO_LOG_TRANSFORM
    FEATURES_TO_LOG_TRANSFORM = to_log_transform(df)
    df[FEATURES_TO_LOG_TRANSFORM] = 10 ** df[FEATURES_TO_LOG_TRANSFORM]
    return df

def compare_config(config , loaded_config):
    """
    Assset that the loaded config from read_graphs is the same as the provided config.
    Especially for the input_features and the max_events, tel_ids. 
    These features if different could produre different graphs


    Parameters
    ----------
    config : dict
        The configuration dictionary used to create the graphs
    loaded_config : dict
        The configuration dictionary loaded from the file.

    Returns
    -------
    bool
        True if the loaded config is the same as the provided config, False otherwise.
    """
    relevant_config_keys = ["data"]
    for part in relevant_config_keys:
        if config.get(part) != loaded_config.get(part):
            return False
    return True




def get_graphs_from_file(filename, config):
    """
    Get graphs from a file.
    If the file is a HDF5 file, the graphs are produced from the table data.
    If the file is a pickle file, the graphs are loaded from the file.

    Parameters
    ----------
    filename : str or Path
        The path to the file containing the data.
    config : dict
        The configuration dictionary.
    """
    filename = Path(filename)
    if filename.name.endswith(".h5"):
        data = read_table(filename, tel_ids=config["data"]["tel_ids"], stop=config["data"]["max_events"])
        data = preprocess_data(data, training_features=config["data"]["input_features"])
        graph_data_list = make_graphs(data, config["data"]["input_features"])
    elif filename.name.endswith(".graphs"):
        graph_data_list, loaded_config = read_graphs(filename)
        # assert that the loaded config from read_graphs is the same as the provided config
        if not compare_config(config, loaded_config):
            raise ValueError("Loaded config is not the same as the provided config")

    else:
        raise ValueError("Invalid file format. Must be either a HDF5 file or a pickle file.")

    return graph_data_list


def read_table(filename, tel_ids=None, stop=None):
    """
    Read data from a file and return it as a pandas DataFrame.

    Parameters
    ----------
    filename : str
        The path to the input file.
    tel_ids : list or None, optional
        A list of telescope IDs to include in the data. If None, all telescopes are included. Default is None.
    stop : int or None, optional
        The number of events to read from the file. If None, all events are read. Default is None.
    Returns
    -------
    pandas.DataFrame
        The data read from the file as a pandas DataFrame.
    """
    loader = TableLoader(
        input_url=filename,
        load_dl1_parameters=True,
        load_dl2=True,
        load_instrument=True,
        load_simulated=True,
        load_true_parameters=True,
    )
    data = loader.read_telescope_events(telescopes=tel_ids, stop=stop)
    data.remove_columns([col for col in data.colnames if len(data[col].shape) > 1])
    return data.to_pandas()


def read_graphs(filename):
    """
    Read graphs from a file.

    Parameters
    ----------
    filename : str or Path
        The path to the file containing the graphs.
    Returns
    -------
    graphs : object
        The graphs loaded from the file.
    """
    graphs, config = torch.load(filename)
    return graphs, config


def save_graphs(graphs, config, filename, overwrite=False):
    """
    Save graphs to a file.

    Parameters
    ----------
    graphs : object
        The graphs to save.
    filename : str or Path
        The path to the file where the graphs will be saved.
    """
    if not overwrite and Path(filename).exists():
        raise FileExistsError(f"File {filename} already exists. Use --overwrite to overwrite the file.")
    torch.save([graphs, config], filename)


def filter_events(events, filters=None, finite_params=None):
    
    from astropy.table import Table

    if isinstance(events, Table):
        events_df = events.to_pandas()
    else:
        events_df = events

    filter = np.ones(len(events_df), dtype=bool)
    filters = {} if filters is None else filters

    for col, (lower_limit, upper_limit) in filters.items():
        filter &= (events_df[col] >= lower_limit) & (events_df[col] <= upper_limit)

    if finite_params is not None:
        _finite_params = list(set(finite_params).intersection(list(events_df.columns)))
        events_df[_finite_params] = events_df[_finite_params].replace([np.inf, -np.inf], np.nan)
        not_finite_mask = events_df[_finite_params].isna()
        filter &= ~(not_finite_mask.any(axis=1))
        not_finite_counts = (not_finite_mask).sum(axis=0)[_finite_params]
        if (not_finite_counts > 0).any():
            warnings.warn("Data contains not-predictable events.")
            warnings.warn("Column | Number of non finite values")
            for k, v in not_finite_counts.items():
                if v > 0:
                    warnings.warn(f"{k} : {v}")

    # if pandas DataFrame, return pandas DataFrame
    filter = filter.to_numpy() if hasattr(filter, "to_numpy") else filter
    events = events[filter]

    return events


TEL_TYPE_MAPPER = {
    "LST": 0,
    "MST": 1,
    "SST": 2,
}
SHOWER_ID_MAPPER = {101: 1, 0: 0}


def encode_tel_type(dataframe):
    """
    Encode the telescope type using a LabelEncoder.
    """
    dataframe["type"] = dataframe["type"].map(TEL_TYPE_MAPPER)
    return dataframe


def decode_tel_type(dataframe):
    """
    Decode the telescope type using a LabelEncoder.
    """
    dataframe["type"] = dataframe["type"].map({v: k for k, v in TEL_TYPE_MAPPER.items()})
    return dataframe


def encode_shower_id(dataframe):
    dataframe["true_shower_primary_id"] = dataframe["true_shower_primary_id"].map(SHOWER_ID_MAPPER)
    return dataframe


def decode_shower_id(dataframe):
    inverse_mapper = {v: k for k, v in SHOWER_ID_MAPPER.items()}
    dataframe["true_shower_primary_id"] = dataframe["true_shower_primary_id"].map(inverse_mapper)
    return dataframe


def preprocess_data(merged_table, training_features):
    """
    Preprocesses the data by applying filters, transformations, normalization, and encoding.

        Parameters
        ----------
        merged_table : pandas.DataFrame
            The merged table containing the data.
        training_features : list
            The list of features used for training the RF models.

        Returns
        -------
        pandas.DataFrame
            The preprocessed data.
    """
    # These filters are the ones used to train the RF models
    filters = {
        "hillas_intensity": [50, np.inf],
        "hillas_width": [0, np.inf],
        "morphology_n_pixels": [3, np.inf],
        "leakage_intensity_width_2": [0, 0.5],
    }

    merged_table = filter_events(merged_table, filters=filters, finite_params=training_features)

    log_transform(merged_table)
    encode_tel_type(merged_table)
    if "true_shower_primary_id" in merged_table.columns:
        encode_shower_id(merged_table)  # encode shower id
    # drop na for columns that contains only nan
    merged_table.dropna(axis=1, how="all", inplace=True)
    merged_table.dropna(inplace=True)

    return merged_table


def inverse_process_data(dataframe):
    """
    Inverse the preprocessing steps applied to the data.
    """
    inverse_log_transform(dataframe)
    return dataframe


def get_features_df(dataloader, task_type):
    """
    Convert a dataloader to a pandas DataFrame containing the features.

    Parameters:
    -----------
    dataloader : DataLoader
        The dataloader containing the dataset.
    task_type : str
        The type of task, either "classification", "energy", "direction", or "multitask".

    Returns:
    --------
    pd.DataFrame
        A pandas DataFrame containing the features.
    """
    data = []
    for data_point in dataloader.dataset:
        data_dict = {
            "obs_id": data_point.obs_id.item(),
            "event_id": data_point.event_id.item(),
            "true_energy": 10 ** data_point.true_energy.item(),
            "true_shower_id": data_point.particle_type.item(),
            "true_alt": data_point.true_alt.item(),
            "true_az": data_point.true_az.item(),
        }

        if task_type in ["classification", "multitask"]:
            data_dict["RandomForest_class_reco"] = data_point.rf_reco_class.item()

        if task_type in ["direction", "multitask"]:
            data_dict["RandomForest_reco_alt"] = data_point.rf_reco_direction[0].item()
            data_dict["RandomForest_reco_az"] = data_point.rf_reco_direction[1].item()

        if task_type in ["energy", "multitask"]:
            data_dict["RandomForest_reco_energy"] = data_point.rf_reco.item()

        data.append(data_dict)

    df = pd.DataFrame(data)
    return df



def add_prediction_to_dataframe(df, predictions, task_type):



    """
    Add predictions to a DataFrame based on the task type.

        Parameters
        ----------
        df : pandas.DataFrame
            The DataFrame to which predictions will be added.
        predictions : list
            A list of predictions.
        task_type : str
            The type of task, either "classification", "regression", "direction_regression", "multitask", or "multignn".

        Returns
        -------
        pandas.DataFrame
            The DataFrame with predictions added.

        Raises
        ------
        ValueError
            If an invalid task type is provided.

    """

    if task_type == "classification":
        class_predictions = np.concatenate(predictions)
        df["Graph_reco_gammaness"] = class_predictions

    elif task_type == "energy":
        log_energy_predictions = np.concatenate([p.numpy() for p in predictions])
        energy_predictions = 10**log_energy_predictions
        df["Graph_reco_energy"] = energy_predictions

    elif task_type == "direction":
        alt_predictions = np.concatenate([p[:, 0].numpy() for p in predictions])
        az_predictions = np.concatenate([p[:, 1].numpy() for p in predictions])
        df["Graph_reco_alt"] = alt_predictions
        df["Graph_reco_az"] = az_predictions

    else:
        raise ValueError(
            "Invalid task type. Must be one of 'classification', 'regression', 'direction_regression', 'multitask', or 'multignn'."
        )

    return df
