import os
import shutil
import zipfile
from pathlib import Path
import gradio as gr

from lcmv_class import LCMVSourceEstimator

# Built-in files (must be in repo root)
BUILT_IN_GPSC = Path("ghw280_from_egig.gpsc")
BUILT_IN_SRC = Path("fsaverage-vol-5mm-src.fif")

if not BUILT_IN_GPSC.is_file():
    raise FileNotFoundError(f"Required montage file not found: {BUILT_IN_GPSC}")
if not BUILT_IN_SRC.is_file():
    raise FileNotFoundError(f"Required source space file not found: {BUILT_IN_SRC}")

# Prepare fsaverage structure
FS_DIR = Path("derivatives/lcmv/fsaverage")
BEM_DIR = FS_DIR / "bem"
BEM_DIR.mkdir(parents=True, exist_ok=True)

EXPECTED_SRC = Path("derivatives/lcmv/fsaverage-vol-5mm-src.fif")
EXPECTED_SRC.parent.mkdir(parents=True, exist_ok=True)
if not EXPECTED_SRC.exists():
    shutil.copy(BUILT_IN_SRC, EXPECTED_SRC)

# Fixed subject/task
SUBJECT_ID = "sub"
TASK = ""

def run_lcmv(
    cleaned_fif,
    run_difumo: bool = True,
    reg: float = 0.01,
    n_jobs: int = 1,
    download_option: str = "full_zip",  # "full_zip" or "difumo_only"
):
    abs_base = Path(".").resolve()
    subject_output = abs_base / "derivatives" / "lcmv" / f"{SUBJECT_ID}_{TASK}"
    if subject_output.exists():
        shutil.rmtree(subject_output)
    subject_output.mkdir(parents=True)

    # Copy input
    fif_path = subject_output / "cleaned_input.fif"
    shutil.copy(cleaned_fif.name, fif_path)

    # Config
    config = {
        'project_base': str(abs_base),
        'subject_id': SUBJECT_ID,
        'task': TASK,
        'ica_file_path': str(fif_path.relative_to(abs_base)),
        'gpsc_file_path': str(abs_base / BUILT_IN_GPSC),
        'reg': reg,
        'n_jobs': n_jobs,
    }

    try:
        os.environ['SUBJECTS_DIR'] = str(abs_base / "derivatives" / "lcmv")
        estimator = LCMVSourceEstimator(config)
        metadata = estimator.run_enhanced_computation()

        if run_difumo:
            difumo_config = {'dimension': 512, 'resolution_mm': 2}
            estimator.run_difumo_extraction(difumo_config=difumo_config)

        # ✅ Dual output logic
        if download_option == "difumo_only":
            difumo_file = subject_output / "difumo_time_courses.npy"
            if difumo_file.exists():
                return str(difumo_file)
            else:
                raise FileNotFoundError("DiFuMo file not found. Ensure 'Extract DiFuMo' is enabled.")

        else:  # full_zip
            zip_path = abs_base / f"{SUBJECT_ID}_{TASK}_lcmv_output.zip"
            with zipfile.ZipFile(zip_path, "w", zipfile.ZIP_DEFLATED) as zf:
                for root, _, files in os.walk(subject_output):
                    for file in files:
                        file_path = Path(root) / file
                        arcname = file_path.relative_to(abs_base)
                        zf.write(file_path, arcname)
            return str(zip_path)

    except Exception as e:
        error_log = subject_output / "lcmv_error.log"
        with open(error_log, "w") as f:
            f.write(f"LCMV failed:\n{str(e)}")
        return str(error_log)

# Gradio UI
with gr.Blocks(theme=gr.themes.Base(), title="LCMV Source Estimation") as demo:
    gr.Markdown("# LCMV Source Estimation")
    gr.Markdown("Upload a cleaned EEG file (output from ICA pipeline). Uses built-in `ghw280_from_egig.gpsc` and precomputed source space.")

    with gr.Row():
        with gr.Column():
            fif_input = gr.File(label="Cleaned EEG (.fif)", file_types=[".fif"])
            run_difumo = gr.Checkbox(True, label="Extract DiFuMo 512 time courses")
            reg = gr.Number(0.01, label="LCMV Regularization (reg)")
            n_jobs = gr.Number(1, label="Parallel Jobs (n_jobs)", precision=0)
            download_option = gr.Radio(
                choices=["full_zip", "difumo_only"],
                value="full_zip",
                label="Download Option",
                info="• full_zip: All outputs (STC, metadata, plots, DiFuMo)\n• difumo_only: Only difumo_time_courses.npy (for connectivity)"
            )
            run_btn = gr.Button("Run LCMV", variant="primary")

        with gr.Column():
            output_file = gr.File(label="Download Output")

    run_btn.click(
        fn=run_lcmv,
        inputs=[fif_input, run_difumo, reg, n_jobs, download_option],
        outputs=output_file,
    )

demo.launch()