import torch
import networkx as nx
from torch_geometric.data import Data
from tqdm import tqdm
import numpy as np
from itertools import combinations

import warnings
warnings.filterwarnings("ignore")


def make_graphs(merged_table, input_features, min_tel=1):
    """
    Create a list of graphs from the merged_table DataFrame.
    
    Parameters:
    merged_table (DataFrame): A pandas DataFrame containing the merged and transformed data.
    input_features (list): A list of strings containing the names of the input features.
    min_tel (int): The minimum number of telescopes required to create a graph.
    
    Returns:
    graph_data_list (list): A list of Data objects containing the graph data and the corresponding true values.
    """

    # Group data by 'obs_id' and 'event_id'
    grouped_data = merged_table.groupby(['obs_id', 'event_id', 'true_shower_primary_id'])

    graph_data_list = []
    print("Producing graphs...")
    for (obs_id, event_id, particle_type), group_df in tqdm(grouped_data):
        # Precompute all needed values
        energy = group_df['true_energy'].iloc[0]
        true_alt = group_df['true_alt'].iloc[0] if 'true_alt' in group_df.columns else 0
        true_az = group_df['true_az'].iloc[0] if 'true_az' in group_df.columns else 0

        energie_reco = group_df['RandomForestRegressor_energy'].iloc[0] if 'RandomForestRegressor_energy' in group_df.columns else 0
        class_reco = group_df['RandomForestClassifier_prediction'].iloc[0] if 'RandomForestClassifier_prediction' in group_df.columns else None
        particle_type = int(group_df['true_shower_primary_id'].iloc[0]) if 'true_shower_primary_id' in group_df.columns else None
        alt_reco = group_df['HillasReconstructor_alt'].iloc[0] if 'HillasReconstructor_alt' in group_df.columns else None
        az_reco = group_df['HillasReconstructor_az'].iloc[0] if 'HillasReconstructor_az' in group_df.columns else None

        # Create a dictionary of positions indexed by tel_id
        pos_dict = group_df.set_index('tel_id')[['pos_x', 'pos_y', 'pos_z']].to_dict('index')
        
        # Create a dictionary to map tel_id to an index
        tel_id_to_index = {tel_id: idx for idx, tel_id in enumerate(group_df['tel_id'])}
        index_to_tel_id = {idx: tel_id for tel_id, idx in tel_id_to_index.items()}

        # Initialize the graph and add nodes
        G = nx.Graph()
        for idx in tel_id_to_index.values():
            G.add_node(idx)
        
        # Generate edges and calculate distances
        edge_index_list = []
        edge_attr_list = []
        for idx1, idx2 in combinations(tel_id_to_index.values(), 2):
            tel1 = index_to_tel_id[idx1]
            tel2 = index_to_tel_id[idx2]
            pos1 = torch.tensor(list(pos_dict[tel1].values()))
            pos2 = torch.tensor(list(pos_dict[tel2].values()))
            distance = torch.norm(pos1 - pos2).item()

            edge_index_list.append([idx1, idx2])
            edge_attr_list.append([distance])

        if len(G.nodes) >= min_tel:
            # Convert the networkx graph to a PyTorch Geometric Data object
            edge_index = torch.tensor(edge_index_list, dtype=torch.long).t().contiguous()
            x = torch.tensor(group_df[input_features].values, dtype=torch.float)
            edge_attr = torch.tensor(edge_attr_list, dtype=torch.float)
            data = Data(x=x, edge_index=edge_index, edge_attr=edge_attr)
            
            data.true_alt = torch.tensor([true_alt], dtype=torch.float)
            data.true_az = torch.tensor([true_az], dtype=torch.float)
            data.direction  = torch.stack([data.true_alt, data.true_az], dim=1)  # Necessary ?
            data.true_energy = torch.tensor([energy], dtype=torch.float)
            data.rf_reco = torch.tensor(energie_reco, dtype=torch.float)  # Global RF reconstructed energies
            data.obs_id = torch.tensor(obs_id, dtype=torch.int)
            data.event_id = torch.tensor(event_id, dtype=torch.int)
            
            # Handle optional fields
            if alt_reco is not None and az_reco is not None:
                data.rf_reco_direction = torch.tensor([alt_reco, az_reco], dtype=torch.float)  # RF reconstructed direction
            else:
                data.rf_reco_direction = torch.tensor([np.nan, np.nan], dtype=torch.float)

            if particle_type is not None:
                data.particle_type = torch.tensor(particle_type, dtype=torch.long)  # true particle class (label)

            if class_reco is not None:
                data.rf_reco_class = torch.tensor(class_reco, dtype=torch.float)  # RF reconstructed class
            else:
                data.rf_reco_class = torch.tensor(np.nan, dtype=torch.float) 
            
            graph_data_list.append(data)

    return graph_data_list  # return list of Data objects
