import os
import pickle

import pandas as pd
from tqdm import tqdm

import torch
from torch_geometric.data import Data, InMemoryDataset
import torch_geometric.utils as geoutils

from rdkit import Chem, RDLogger



def label2onehot(labels, dim, device=None):
    """Convert label indices to one-hot vectors."""
    out = torch.zeros(list(labels.size())+[dim])
    if device:
        out = out.to(device)

    out.scatter_(len(out.size())-1,labels.unsqueeze(-1),1.)

    return out.float()


def get_encoders_decoders(raw_file1, raw_file2, max_atom):
    """
    Given two raw SMILES files, either load the atom and bond encoders/decoders
    if they exist (naming them based on the file names) or create and save them.

    Parameters:
        raw_file1 (str): Path to the first SMILES file.
        raw_file2 (str): Path to the second SMILES file.
        max_atom (int): Maximum allowed number of atoms in a molecule.

    Returns:
        atom_encoder (dict): Mapping from atomic numbers to indices.
        atom_decoder (dict): Mapping from indices to atomic numbers.
        bond_encoder (dict): Mapping from bond types to indices.
        bond_decoder (dict): Mapping from indices to bond types.
    """
    # Determine unique suffix based on the two file names (alphabetically sorted for consistency)
    name1 = os.path.splitext(os.path.basename(raw_file1))[0]
    name2 = os.path.splitext(os.path.basename(raw_file2))[0]
    sorted_names = sorted([name1, name2])
    suffix = f"{sorted_names[0]}_{sorted_names[1]}"

    # Define encoder/decoder directories and file paths
    enc_dir = os.path.join("data", "encoders")
    dec_dir = os.path.join("data", "decoders")
    atom_encoder_path = os.path.join(enc_dir, f"atom_{suffix}.pkl")
    atom_decoder_path = os.path.join(dec_dir, f"atom_{suffix}.pkl")
    bond_encoder_path = os.path.join(enc_dir, f"bond_{suffix}.pkl")
    bond_decoder_path = os.path.join(dec_dir, f"bond_{suffix}.pkl")

    # If all files exist, load and return them
    if (os.path.exists(atom_encoder_path) and os.path.exists(atom_decoder_path) and 
        os.path.exists(bond_encoder_path) and os.path.exists(bond_decoder_path)):
        with open(atom_encoder_path, "rb") as f:
            atom_encoder = pickle.load(f)
        with open(atom_decoder_path, "rb") as f:
            atom_decoder = pickle.load(f)
        with open(bond_encoder_path, "rb") as f:
            bond_encoder = pickle.load(f)
        with open(bond_decoder_path, "rb") as f:
            bond_decoder = pickle.load(f)
        print("Loaded existing encoders/decoders!")
        return atom_encoder, atom_decoder, bond_encoder, bond_decoder

    # Otherwise, create the encoders/decoders
    print("Creating new encoders/decoders...")
    # Read SMILES from both files (assuming one SMILES per row, no header)
    smiles1 = pd.read_csv(raw_file1, header=None)[0].tolist()
    smiles2 = pd.read_csv(raw_file2, header=None)[0].tolist()
    smiles_combined = smiles1 + smiles2

    atom_labels = set()
    bond_labels = set()
    max_length = 0
    filtered_smiles = []
    
    # Process each SMILES: keep only valid molecules with <= max_atom atoms
    for smiles in tqdm(smiles_combined, desc="Processing SMILES"):
        mol = Chem.MolFromSmiles(smiles)
        if mol is None:
            continue
        molecule_size = mol.GetNumAtoms()
        if molecule_size > max_atom:
            continue
        filtered_smiles.append(smiles)
        # Collect atomic numbers
        atom_labels.update([atom.GetAtomicNum() for atom in mol.GetAtoms()])
        max_length = max(max_length, molecule_size)
        # Collect bond types
        bond_labels.update([bond.GetBondType() for bond in mol.GetBonds()])
    
    # Add a PAD symbol (here using 0 for atoms)
    atom_labels.add(0)
    atom_labels = sorted(atom_labels)
    
    # For bonds, prepend the PAD bond type (using rdkit's BondType.ZERO)
    bond_labels = sorted(bond_labels)
    bond_labels = [Chem.rdchem.BondType.ZERO] + bond_labels

    # Create encoder and decoder dictionaries
    atom_encoder = {l: i for i, l in enumerate(atom_labels)}
    atom_decoder = {i: l for i, l in enumerate(atom_labels)}
    bond_encoder = {l: i for i, l in enumerate(bond_labels)}
    bond_decoder = {i: l for i, l in enumerate(bond_labels)}

    # Ensure directories exist
    os.makedirs(enc_dir, exist_ok=True)
    os.makedirs(dec_dir, exist_ok=True)

    # Save the encoders/decoders to disk
    with open(atom_encoder_path, "wb") as f:
        pickle.dump(atom_encoder, f)
    with open(atom_decoder_path, "wb") as f:
        pickle.dump(atom_decoder, f)
    with open(bond_encoder_path, "wb") as f:
        pickle.dump(bond_encoder, f)
    with open(bond_decoder_path, "wb") as f:
        pickle.dump(bond_decoder, f)

    print("Encoders/decoders created and saved.")
    return atom_encoder, atom_decoder, bond_encoder, bond_decoder

def load_molecules(data=None, b_dim=32, m_dim=32, device=None, batch_size=32):
    data = data.to(device)
    a = geoutils.to_dense_adj(
        edge_index = data.edge_index,
        batch=data.batch,
        edge_attr=data.edge_attr,
        max_num_nodes=int(data.batch.shape[0]/batch_size)
    )
    x_tensor = data.x.view(batch_size,int(data.batch.shape[0]/batch_size),-1)
    a_tensor = label2onehot(a, b_dim, device)

    a_tensor_vec = a_tensor.reshape(batch_size,-1)
    x_tensor_vec = x_tensor.reshape(batch_size,-1)
    real_graphs = torch.concat((x_tensor_vec,a_tensor_vec),dim=-1)

    return real_graphs, a_tensor, x_tensor