import gradio as gr
import os
import shutil
import subprocess
from llama_cpp import Llama
from pymatgen.core import Structure
from huggingface_hub import hf_hub_download

# --------------------
# Paths & Config
# --------------------
GEN_DIR = "generation_cifs"

REPO_ID = "andreamena/crystext-gguf"
FILENAME = "crystext-rl.gguf"

# Fixed evaluation compositions (from your dataset)
KNOWN_MATERIALS = [
    "NaTiVS4",
    "ScAlRh2",
    "Fe3Te2Se",
    "Eu2CrSbO6",
    "Fe(HO)2",
    "Sm(CoB)2",
    "GaTe",
    "Zr2VRe3",
    "LiMnO2",
    "AlAu4",
]

# --------------------
# Load GGUF model
# --------------------
print("⏳ Downloading Model from Hub... this may take a minute...")
try:
    model_path = hf_hub_download(repo_id=REPO_ID, filename=FILENAME)
    print(f"✅ Model downloaded to: {model_path}")

    print("⏳ Initializing CPU Engine...")
    llm = Llama(
        model_path=model_path,
        n_ctx=2048,
        n_threads=2,      # HF free tier = 2 vCPUs
        verbose=True
    )
    print("✅ Engine Ready!")
except Exception as e:
    print(f"❌ Error setting up model: {e}")
    llm = None

# --------------------
# Generation helpers
# --------------------
def generate_single_cif(composition, space_group, temperature):
    prompt = f"""Below is an instruction that describes a task, paired with an input that provides further context.

### Instruction:
Generate CIF for the given material description

### Input:
Material composition is {composition}. It has a space group number {space_group}.

### Response:
"""
    output = llm(
        prompt,
        max_tokens=2048,
        temperature=temperature,
        stop=["</s>"],
        echo=False
    )
    return output["choices"][0]["text"]

# --------------------
# Main pipeline
# --------------------
def run_demo(composition, space_group, num_samples, temperature, progress=gr.Progress()):
    shutil.rmtree(GEN_DIR, ignore_errors=True)
    os.makedirs(GEN_DIR, exist_ok=True)

    if llm is None:
        return None, "Model failed to load.", "N/A"

    status_msg = f"Generating {num_samples} candidates (CPU: ~6 min/sample)\n"
    generated_files = []

    for i in progress.tqdm(range(int(num_samples)), desc="Generating structures"):
        try:
            cif_text = generate_single_cif(composition, space_group, temperature)
            cif_text = cif_text.replace("```cif", "").replace("```", "").strip()

            path = os.path.join(GEN_DIR, f"sample_{i}.cif")
            with open(path, "w") as f:
                f.write(cif_text)

            generated_files.append(path)
        except Exception as e:
            print(f"Sample {i} failed: {e}")

    if not generated_files:
        return None, "No valid CIFs generated.", "N/A"

    # Evaluation
    progress(0.9, desc="Evaluating structures...")
    try:
        subprocess.run(["python", "select_best_cif.py"], check=True)
        eval_text = subprocess.check_output(["python", "evaluation.py"], text=True)

        final_path = "final.cif"
        if os.path.exists(final_path):
            struct = Structure.from_file(final_path)
            sg_info = struct.get_space_group_info(symprec=0.1)
            return final_path, f"{status_msg}\n{eval_text}", f"{sg_info}"
        else:
            return None, f"{status_msg}\nEvaluation finished but no final.cif found.", "N/A"

    except Exception as e:
        return None, f"Evaluation error: {e}", "N/A"

# --------------------
# Gradio UI
# --------------------
with gr.Blocks() as demo:
    gr.Markdown("# CrysText-RL: Crystal Generation")
    gr.Markdown(f"Running GGUF model sourced from: `{REPO_ID}`")

    with gr.Row():
        with gr.Column():
            composition = gr.Dropdown(
                choices=KNOWN_MATERIALS,
                value="NaTiVS4",
                label="Composition (evaluation set)"
            )

            space_group = gr.Slider(
                minimum=1,
                maximum=230,
                step=1,
                value=8,
                label="Target Space Group"
            )

            num_samples = gr.Slider(
                1, 6, value=1, step=1,
                label="Samples (CPU-limited)"
            )

            temperature = gr.Slider(
                0.6, 1.0, value=0.8, step=0.05,
                label="Temperature"
            )

            run_btn = gr.Button("Generate Crystal Structure", variant="primary")

        with gr.Column():
            sg_output = gr.Textbox(label="Predicted Space Group")
            eval_output = gr.Textbox(label="Evaluation Metrics", lines=10)
            best_cif = gr.File(label="Download Best CIF")

    run_btn.click(
        run_demo,
        inputs=[composition, space_group, num_samples, temperature],
        outputs=[best_cif, eval_output, sg_output],
    )

demo.launch(
    auth=(os.environ.get("app_user"), os.environ.get("app_pass")),
    ssr_mode=False

)