import spaces

# Standard Libraries
import os
import io
import csv
import json
import glob
import random
import tempfile
import atexit
from datetime import datetime

# Third-Party Libraries
import numpy as np
import pandas as pd
import torch
import imageio
from rdkit import Chem
from rdkit.Chem import Draw
import gradio as gr

# Local Modules
from evaluator import Evaluator
from loader import load_graph_decoder

# --------------------------- Configuration Constants --------------------------- #

DATA_DIR = 'data'
EVALUATORS_DIR = 'evaluators'
FLAGGED_FOLDER = "flagged"
KNOWN_LABELS_FILE = os.path.join(DATA_DIR, 'known_labels.csv')
KNOWN_SMILES_FILE = os.path.join(DATA_DIR, 'known_polymers.csv')

ALL_PROPERTIES = ['CH4', 'CO2', 'H2', 'N2', 'O2']
MODEL_NAME_MAPPING = {
    "model_all": "Graph DiT (trained on labeled + unlabeled)",
    "model_labeled": "Graph DiT (trained on labeled)"
}

GIF_TEMP_PREFIX = "polymer_gifs_"

# --------------------------- Data Loading --------------------------- #

def load_known_data():
    """Load known labels and SMILES data from CSV files."""
    try:
        known_labels = pd.read_csv(KNOWN_LABELS_FILE)
        known_smiles = pd.read_csv(KNOWN_SMILES_FILE)
        return known_labels, known_smiles
    except Exception as e:
        raise FileNotFoundError(f"Error loading data files: {e}")

# Load data
known_labels, known_smiles = load_known_data()

# --------------------------- Evaluator Setup --------------------------- #

def initialize_evaluators(properties, evaluators_dir):
    """Initialize evaluators for each property."""
    evaluators = {}
    for prop in properties:
        evaluator_path = os.path.join(evaluators_dir, f'{prop}.joblib')
        evaluators[prop] = Evaluator(evaluator_path, prop)
    return evaluators

evaluators = initialize_evaluators(ALL_PROPERTIES, EVALUATORS_DIR)

# --------------------------- Property Ranges --------------------------- #

def get_property_ranges(labels, properties):
    """Get min and max values for each property."""
    return {prop: (labels[prop].min(), labels[prop].max()) for prop in properties}

property_ranges = get_property_ranges(known_labels, ALL_PROPERTIES)

# --------------------------- Temporary Directory Setup --------------------------- #

temp_dir = tempfile.mkdtemp(prefix=GIF_TEMP_PREFIX)

def cleanup_temp_files():
    """Clean up temporary GIF files on exit."""
    try:
        for file in glob.glob(os.path.join(temp_dir, "*.gif")):
            os.remove(file)
        os.rmdir(temp_dir)
    except Exception as e:
        print(f"Error during cleanup: {e}")

atexit.register(cleanup_temp_files)

# --------------------------- Utility Functions --------------------------- #

def random_properties():
    """Select a random set of properties from known labels."""
    return known_labels[ALL_PROPERTIES].sample(1).values.tolist()[0]

def load_model(model_choice):
    """Load the graph decoder model based on the choice."""
    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
    model = load_graph_decoder(path=model_choice)
    return model, device

def save_interesting_log(smiles, properties, suggested_properties):
    """Save interesting polymer data to a CSV log file."""
    log_file = os.path.join(FLAGGED_FOLDER, "log.csv")
    os.makedirs(FLAGGED_FOLDER, exist_ok=True)
    file_exists = os.path.isfile(log_file)

    fieldnames = ['timestamp', 'smiles'] + ALL_PROPERTIES + [f'suggested_{prop}' for prop in ALL_PROPERTIES]

    try:
        with open(log_file, 'a', newline='') as csvfile:
            writer = csv.DictWriter(csvfile, fieldnames=fieldnames)
            if not file_exists:
                writer.writeheader()

            log_data = {
                'timestamp': datetime.now().strftime("%Y-%m-%d %H:%M:%S"),
                'smiles': smiles,
                **{prop: value for prop, value in zip(ALL_PROPERTIES, properties)},
                **{f'suggested_{prop}': value for prop, value in suggested_properties.items()}
            }
            writer.writerow(log_data)
    except Exception as e:
        print(f"Error saving log: {e}")

def is_nan_like(x):
    """Check if a value should be treated as NaN."""
    return x == 0 or x == '' or (isinstance(x, float) and np.isnan(x))

def numpy_to_python(obj):
    """Convert NumPy objects to native Python types."""
    if isinstance(obj, np.integer):
        return int(obj)
    elif isinstance(obj, np.floating):
        return float(obj)
    elif isinstance(obj, np.ndarray):
        return obj.tolist()
    elif isinstance(obj, list):
        return [numpy_to_python(item) for item in obj]
    elif isinstance(obj, dict):
        return {k: numpy_to_python(v) for k, v in obj.items()}
    else:
        return obj

# --------------------------- Graph Generation Function --------------------------- #

@spaces.GPU
def generate_graph(CH4, CO2, H2, N2, O2, guidance_scale, num_nodes, repeating_time, model_state, num_chain_steps, fps):
    """
    Generate a polymer graph based on the input properties and model.
    Returns generation results including SMILES, images, and properties.
    """
    print('Generating graph...')
    model, device = model_state
    properties = [CH4, CO2, H2, N2, O2]

    # Handle NaN-like values
    properties = [None if is_nan_like(prop) else prop for prop in properties]

    nan_gases = [gas for gas, prop in zip(ALL_PROPERTIES, properties) if prop is None]
    nan_message = "The following gas properties were treated as NaN: " + (", ".join(nan_gases) if nan_gases else "None")

    num_nodes = None if num_nodes == 0 else num_nodes

    for attempt in range(repeating_time):
        try:
            generated_molecule, img_list = model.generate(
                properties,
                device=device,
                guide_scale=guidance_scale,
                num_nodes=num_nodes,
                number_chain_steps=num_chain_steps
            )

            gif_path = None
            if img_list:
                imgs = [np.array(pil_img) for pil_img in img_list]
                imgs.extend([imgs[-1]] * 10)  # Extend the last image for GIF
                gif_path = os.path.join(temp_dir, f"polymer_gen_{random.randint(0, 999999)}.gif")
                imageio.mimsave(gif_path, imgs, format='GIF', fps=fps, loop=0)

            if generated_molecule:
                mol = Chem.MolFromSmiles(generated_molecule)
                if mol:
                    standardized_smiles = Chem.MolToSmiles(mol, isomericSmiles=True)
                    is_novel = standardized_smiles not in known_smiles['SMILES'].values
                    novelty_status = "Novel (Not in Labeled Set)" if is_novel else "Not Novel (Exists in Labeled Set)"
                    img = Draw.MolToImage(mol)

                    # Evaluate the generated molecule
                    suggested_properties = {prop: evaluator([standardized_smiles])[0] for prop, evaluator in evaluators.items()}

                    suggested_properties_text = "\n".join([f"**Suggested {prop}:** {value:.2f}" for prop, value in suggested_properties.items()])

                    return (
                        f"**Generated polymer SMILES:** `{standardized_smiles}`\n\n"
                        f"**{nan_message}**\n\n"
                        f"**{novelty_status}**\n\n"
                        f"**Suggested Properties:**\n{suggested_properties_text}",
                        img,
                        gif_path,
                        standardized_smiles,
                        properties,
                        suggested_properties
                    )
        except Exception as e:
            print(f"Attempt {attempt + 1} failed: {e}")
            continue

    # If all attempts fail
    return (
        f"**Generation failed:** Could not generate a valid molecule after {repeating_time} attempts.\n\n**{nan_message}**",
        None,
        None,
        "",
        [],
        {}
    )

# --------------------------- Feedback Processing --------------------------- #

def process_feedback(checkbox_value, smiles, properties, suggested_properties):
    """
    Process user feedback. If the user finds the polymer interesting,
    log it accordingly.
    """
    if checkbox_value and smiles:
        save_interesting_log(smiles, properties, suggested_properties)
        return "Thank you for your feedback! This polymer has been saved to our interesting polymers log."
    return "Thank you for your feedback!"

# --------------------------- Model Switching --------------------------- #

def switch_model(choice):
    """Switch the model based on user selection."""
    internal_name = next(key for key, value in MODEL_NAME_MAPPING.items() if value == choice)
    return load_model(internal_name)

# --------------------------- Gradio Interface Setup --------------------------- #

def create_gradio_interface():
    """Create and return the Gradio Blocks interface."""
    with gr.Blocks(title="Polymer Design with GraphDiT") as iface:
        # Navigation Bar
        with gr.Row(elem_id="navbar"):
            gr.Markdown("""
            <div style="text-align: center;">
                <h1>🔗🔬 Polymer Design with GraphDiT</h1>
                <div style="display: flex; gap: 20px; justify-content: center; align-items: center; margin-top: 10px;">
                    <a href="https://github.com/liugangcode/Graph-DiT" target="_blank" style="display: flex; align-items: center; gap: 5px; text-decoration: none; color: inherit;">
                        <img src="https://img.icons8.com/ios-glyphs/30/000000/github.png" alt="GitHub" />
                        <span>View Code</span>
                    </a>
                    <a href="https://arxiv.org/abs/2401.13858" target="_blank" style="text-decoration: none; color: inherit;">
                        📄 View Paper
                    </a>
                </div>
            </div>
            """)

        # Main Description
        gr.Markdown("""
        ## Introduction

        Input the desired gas barrier properties for CH₄, CO₂, H₂, N₂, and O₂ to generate novel polymer structures. The results are visualized as molecular graphs and represented by SMILES strings if they are successfully generated. **Note:** Gas barrier values set to 0 will be treated as `NaN` (unconditionally). If the generation fails, please retry or increase the number of repetition attempts.
        """)

        # Model Selection
        model_choice = gr.Radio(
            choices=list(MODEL_NAME_MAPPING.values()),
            label="Model Zoo",
            value=MODEL_NAME_MAPPING["model_labeled"]
        )

        # Model Description Accordion
        with gr.Accordion("🔍 Model Description", open=False):
            gr.Markdown("""
            ### GraphDiT: Graph Diffusion Transformer

            GraphDiT is a graph diffusion model designed for targeted molecular generation. It employs a conditional diffusion process to iteratively refine molecular structures based on user-specified properties.

            We have collected a labeled polymer database for gas permeability from [Membrane Database](https://research.csiro.au/virtualscreening/membrane-database-polymer-gas-separation-membranes/). Additionally, we utilize unlabeled polymer structures from [PolyInfo](https://polymer.nims.go.jp/).

            The gas permeability ranges from 0 to over ten thousand, with only hundreds of labeled data points, making this task particularly challenging.

            We are actively working on improving the model. We welcome any feedback regarding model usage or suggestions for improvement.

            #### Currently, we have two variants of Graph DiT:
            - **Graph DiT (trained on labeled + unlabeled)**: This model uses both labeled and unlabeled data for training, potentially leading to more diverse/novel polymer generation.
            - **Graph DiT (trained on labeled)**: This model is trained exclusively on labeled data, which may result in higher validity but potentially less diverse/novel outputs.
            """)

        # Citation Accordion
        with gr.Accordion("📄 Citation", open=False):
            gr.Markdown("""
            If you use this model or interface useful, please cite the following paper:
            ```bibtex
            @article{graphdit2024,
              title={Graph Diffusion Transformers for Multi-Conditional Molecular Generation},
              author={Liu, Gang and Xu, Jiaxin and Luo, Tengfei and Jiang, Meng},
              journal={NeurIPS},
              year={2024},
            }
            ```
            """)

        # Initialize Model State
        model_state = gr.State(load_model("model_labeled"))

        # Property Inputs
        with gr.Row():
            CH4_input = gr.Slider(
                minimum=0,
                maximum=property_ranges['CH4'][1],
                value=2.5,
                label=f"CH₄ (Barrier) [0-{property_ranges['CH4'][1]:.1f}]"
            )
            CO2_input = gr.Slider(
                minimum=0,
                maximum=property_ranges['CO2'][1],
                value=15.4,
                label=f"CO₂ (Barrier) [0-{property_ranges['CO2'][1]:.1f}]"
            )
            H2_input = gr.Slider(
                minimum=0,
                maximum=property_ranges['H2'][1],
                value=21.0,
                label=f"H₂ (Barrier) [0-{property_ranges['H2'][1]:.1f}]"
            )
            N2_input = gr.Slider(
                minimum=0,
                maximum=property_ranges['N2'][1],
                value=1.5,
                label=f"N₂ (Barrier) [0-{property_ranges['N2'][1]:.1f}]"
            )
            O2_input = gr.Slider(
                minimum=0,
                maximum=property_ranges['O2'][1],
                value=2.8,
                label=f"O₂ (Barrier) [0-{property_ranges['O2'][1]:.1f}]"
            )

        # Generation Parameters
        with gr.Row():
            guidance_scale = gr.Slider(
                minimum=1,
                maximum=3,
                value=2,
                label="Guidance Scale from Properties"
            )
            num_nodes = gr.Slider(
                minimum=0,
                maximum=50,
                step=1,
                value=0,
                label="Number of Nodes (0 for Random, Larger Graphs Take More Time)"
            )
            repeating_time = gr.Slider(
                minimum=1,
                maximum=10,
                step=1,
                value=3,
                label="Repetition Until Success"
            )
            num_chain_steps = gr.Slider(
                minimum=0,
                maximum=499,
                step=1,
                value=50,
                label="Number of Diffusion Steps to Visualize (Larger Numbers Take More Time)"
            )
            fps = gr.Slider(
                minimum=0.25,
                maximum=10,
                step=0.25,
                value=5,
                label="Frames Per Second"
            )

        # Action Buttons
        with gr.Row():
            random_btn = gr.Button("🔀 Randomize Properties (from Labeled Data)")
            generate_btn = gr.Button("🚀 Generate Polymer")

        # Results Display
        with gr.Row():
            result_text = gr.Textbox(label="📝 Generation Result", lines=10)
            result_image = gr.Image(label="Final Molecule Visualization", type="pil")
            result_gif = gr.Image(label="Generation Process Visualization", type="filepath", format="gif")

        # Feedback Section
        with gr.Row():
            feedback_btn = gr.Button("🌟 I think this polymer is interesting!", interactive=False)
            feedback_result = gr.Textbox(label="Feedback Result", visible=False)

        # Hidden Components to Store Generation Data
        hidden_smiles = gr.Textbox(visible=False)
        hidden_properties = gr.JSON(visible=False)
        hidden_suggested_properties = gr.JSON(visible=False)

        # Event Handlers

        # Model Selection Change
        model_choice.change(
            switch_model,
            inputs=[model_choice],
            outputs=[model_state]
        )

        # Randomize Properties Button
        random_btn.click(
            random_properties,
            outputs=[CH4_input, CO2_input, H2_input, N2_input, O2_input]
        )

        # Generate Polymer Button
        generate_btn.click(
            generate_graph,
            inputs=[
                CH4_input, CO2_input, H2_input, N2_input, O2_input,
                guidance_scale, num_nodes, repeating_time,
                model_state, num_chain_steps, fps
            ],
            outputs=[
                result_text, result_image, result_gif,
                hidden_smiles, hidden_properties, hidden_suggested_properties
            ]
        ).then(
            lambda text, img, gif, smiles, props, sugg_props: (
                smiles if text.startswith("**Generated polymer SMILES:**") else "",
                json.dumps(numpy_to_python(props)),
                json.dumps(numpy_to_python(sugg_props)),
                gr.Button(interactive=text.startswith("**Generated polymer SMILES:**"))
            ),
            inputs=[
                result_text, result_image, result_gif,
                hidden_smiles, hidden_properties, hidden_suggested_properties
            ],
            outputs=[hidden_smiles, hidden_properties, hidden_suggested_properties, feedback_btn]
        )

        # Feedback Button Click
        feedback_btn.click(
            process_feedback,
            inputs=[gr.Checkbox(label="Interested?", value=True, visible=False), hidden_smiles, hidden_properties, hidden_suggested_properties],
            outputs=[feedback_result]
        ).then(
            lambda: gr.Button(interactive=False),
            outputs=[feedback_btn]
        )

        # Reset Feedback Button on Input Changes
        for input_component in [CH4_input, CO2_input, H2_input, N2_input, O2_input, random_btn]:
            input_component.change(
                lambda: None,
                outputs=[feedback_btn],
                _js="() => feedback_btn.interactive = false"
            )

    return iface

# --------------------------- Main Execution --------------------------- #

if __name__ == "__main__":
    interface = create_gradio_interface()
    interface.launch(share=False)