"""
RL-Medical / physics-informed landmark detector.

Loads the bundled single-agent DQN checkpoint (trained to find landmark #13 -
the anterior commissure - in brain MRI) and runs it in inference-only
("play") mode on a NIfTI volume the user uploads.

This wraps the existing CLI (`src/DQN.py --task play`) as a subprocess so the
original, unmodified training/inference code is reused as-is.
"""

import ast
import os
import re
import subprocess
import sys
import tempfile
import spaces
import pydicom
import nibabel as nib

import gradio as gr
import matplotlib
matplotlib.use("Agg")
import matplotlib.pyplot as plt
from matplotlib.animation import FuncAnimation, PillowWriter
import numpy as np
import SimpleITK as sitk

ROOT = os.path.dirname(os.path.abspath(__file__))
SRC_DIR = os.path.join(ROOT, "src")
MODEL_PATH = os.path.join(SRC_DIR, "data", "models", "BrainMRI", "SingleAgent.pt")
LANDMARK_ID = 13  # what SingleAgent.pt was trained on
MAX_GIF_FRAMES = 60  # subsample long trajectories so the animation stays quick to build


@spaces.GPU
def run_inference(nifti_path: str):
    """Runs `DQN.py --task play` on a single volume and returns the row
    printed by the logger with the agent's final (x, y, z) voxel position,
    plus the full step-by-step trajectory (from the STEP_LOC log lines)."""

    with tempfile.TemporaryDirectory() as tmp:
        files_list = os.path.join(tmp, "image_files.txt")
        with open(files_list, "w") as f:
            f.write(nifti_path + "\n")

        cmd = [
            sys.executable, "DQN.py",
            "--task", "play",
            "--load", MODEL_PATH,
            "--files", files_list,
            "--file_type", "brain",
            "--landmarks", str(LANDMARK_ID),
            "--model_name", "Network3d",
            "--viz", "0",
        ]
        result = subprocess.run(
            cmd, cwd=SRC_DIR, capture_output=True, text=True, timeout=600
        )

        if result.returncode != 0:
            raise RuntimeError(
                "Inference failed:\n" + result.stdout[-2000:] + "\n" + result.stderr[-2000:]
            )

        # The logger prints each results row as a Python list literal, and
        # evaluator.py prints one "STEP_LOC: <step> <x> <y> <z>" line per
        # step of the episode (agent 0 only).
        row = None
        trajectory = []  # list of (step, x, y, z)
        for line in result.stdout.splitlines():
            line = line.strip()
            if line.startswith("STEP_LOC:"):
                parts = line.split()
                try:
                    step, x, y, z = int(parts[1]), float(parts[2]), float(parts[3]), float(parts[4])
                    trajectory.append((step, x, y, z))
                except (ValueError, IndexError):
                    continue
            elif line.startswith("[") and line.endswith("]"):
                try:
                    parsed = ast.literal_eval(line)
                    if isinstance(parsed, list) and isinstance(parsed[0], int):
                        row = parsed
                except (ValueError, SyntaxError):
                    continue
        if row is None:
            raise RuntimeError("Could not parse model output:\n" + result.stdout[-2000:])

        # row = [episode, filename, agent_x, agent_y, agent_z, landmark_x/N/A, ...]
        x, y, z = row[2], row[3], row[4]
        return int(round(x)), int(round(y)), int(round(z)), trajectory

def is_dicom(path):
    """
    Returns True if the file is a valid DICOM.
    """
    try:
        pydicom.dcmread(path, stop_before_pixels=True)
        return True
    except Exception:
        return False


def is_nifti(path):
    """
    Returns True if the file is a valid NIfTI.
    """
    try:
        nib.load(path)
        return True
    except Exception:
        return False


def make_preview(nifti_path: str, x: int, y: int, z: int):
    """Renders the axial/coronal/sagittal slices through the predicted point."""
    image = sitk.ReadImage(nifti_path)
    array = sitk.GetArrayFromImage(image)  # z, y, x
    zmax, ymax, xmax = array.shape
    x = min(max(x, 0), xmax - 1)
    y = min(max(y, 0), ymax - 1)
    z = min(max(z, 0), zmax - 1)

    fig, axs = plt.subplots(1, 3, figsize=(9, 3.2))
    views = [
        (array[z, :, :], (x, y), "Axial"),
        (array[:, y, :], (x, z), "Coronal"),
        (array[:, :, x], (y, z), "Sagittal"),
    ]
    for ax, (slice_, point, title) in zip(axs, views):
        ax.imshow(slice_, cmap="gray")
        ax.scatter([point[0]], [point[1]], c="red", s=40, marker="+")
        ax.set_title(title)
        ax.axis("off")
    fig.tight_layout()
    return fig


def make_trajectory_gif(nifti_path: str, trajectory, out_path: str, fps: int = 8):
    """Animates the agent's real search path: each frame is the axial slice
    at that step's z-position, styled to match the project's original
    agent-visualization convention (black background, yellow ROI box = what
    the network sees, blue dot = current location, cyan trail = path so far).
    """
    if not trajectory:
        return None

    image = sitk.ReadImage(nifti_path)
    array = sitk.GetArrayFromImage(image)  # z, y, x
    zmax, ymax, xmax = array.shape
    vmin, vmax = np.percentile(array, 2), np.percentile(array, 99)

    if len(trajectory) > MAX_GIF_FRAMES:
        idxs = np.linspace(0, len(trajectory) - 2, MAX_GIF_FRAMES - 1).astype(int)
        frames = [trajectory[i] for i in idxs] + [trajectory[-1]]
    else:
        frames = trajectory

    fig, ax = plt.subplots(figsize=(4.5, 4.5), facecolor="black")
    roi = 40  # half-size of the "what the network sees" ROI box
    xs_trail, ys_trail = [], []

    def render(frame):
        ax.clear()
        ax.set_facecolor("black")
        ax.set_xticks([])
        ax.set_yticks([])
        step, x, y, z = frame
        x = int(min(max(x, 0), xmax - 1))
        y = int(min(max(y, 0), ymax - 1))
        z = int(min(max(z, 0), zmax - 1))
        ax.imshow(array[z, :, :], cmap="gray", vmin=vmin, vmax=vmax)
        xs_trail.append(x)
        ys_trail.append(y)
        ax.plot(xs_trail, ys_trail, "-", color="#7cf6ff", lw=1.6, alpha=0.9)
        ax.add_patch(plt.Rectangle((x - roi, y - roi), 2 * roi, 2 * roi,
                                    fill=False, edgecolor="#ffcc00", lw=1.4))
        ax.plot(x, y, "o", color="#3a7bff", ms=7, mec="w", mew=0.6)
        ax.set_title(f"Step {step}", color="#ffcc00", fontsize=10)
        fig.suptitle("Agent search path", color="white", fontsize=11)

    anim = FuncAnimation(fig, render, frames=frames)
    anim.save(out_path, writer=PillowWriter(fps=fps), dpi=100)
    plt.close(fig)
    return out_path


def prepare_image(file):
    """
    Returns a NIfTI file path for either NIfTI or DICOM input.
    """

    path = file.name

    # ---------- NIfTI ----------
    if is_nifti(path):
        return path

    # ---------- DICOM ----------
    if is_dicom(path):

        image = sitk.ReadImage(path)

        tmp = tempfile.NamedTemporaryFile(
            suffix=".nii.gz",
            delete=False
        )

        sitk.WriteImage(image, tmp.name)

        return tmp.name

    raise gr.Error("Unsupported medical image format.")


def detect(file):

    image_path = prepare_image(file)

    x, y, z, trajectory = run_inference(image_path)

    fig = make_preview(image_path, x, y, z)

    gif_fd, gif_path = tempfile.mkstemp(suffix=".gif")
    os.close(gif_fd)
    gif_path = make_trajectory_gif(image_path, trajectory, gif_path)

    text = f"""
✅ Analysis Complete

Detected Landmark : #{LANDMARK_ID}

Voxel Coordinates

X : {x}
Y : {y}
Z : {z}

Steps taken to converge : {trajectory[-1][0] if trajectory else 'N/A'}
"""

    return fig, text, gif_path


css = """
.gradio-container{
    max-width:1500px !important;
}

h1{
    text-align:center;
}

.section{
    border-radius:15px;
    padding:15px;
    background:#f8fafc;
    border:1px solid #e5e7eb;
}

.footer{
    text-align:center;
    font-size:13px;
    color:gray;
}
"""

with gr.Blocks(
    css=css,
    title="Communicative RL Medical Imaging"
) as demo:

    gr.Markdown(
        """
# 🧠 Communicative Reinforcement Learning Medical Imaging Platform

### Physics-Informed Landmark Detection for Brain MRI

Upload either

- ✅ NIfTI (.nii / .nii.gz)

or

- ✅ DICOM (.dcm)

The AI automatically detects anatomical landmarks.
"""
    )

    with gr.Row():

        with gr.Column(scale=1):

            gr.Markdown("## 📤 Upload")

            file_input = gr.File(
                label="Medical Image"
                )

            detect_btn = gr.Button(
                "🚀 Start AI Analysis",
                variant="primary",
                size="lg"
            )

            clear_btn = gr.ClearButton(
                components=[file_input],
                value="🗑 Clear"
            )

            gr.Markdown("---")

            metadata = gr.Textbox(
                label="📋 Image Information",
                lines=10,
                interactive=False
            )

        with gr.Column(scale=2):

            with gr.Tabs():

                with gr.Tab("🖼 Visualization"):

                    output_image = gr.Plot(
                        label="Landmark Detection"
                    )

                with gr.Tab("🎞 Agent Path"):

                    gr.Markdown(
                        "The agent starts at a random point and steps through "
                        "the volume toward the landmark. This animates its "
                        "actual path, slice by slice, as it converges. The "
                        "yellow box is the region the network is looking at "
                        "around its current position."
                    )

                    path_gif = gr.Image(
                        label="Search trajectory",
                        type="filepath"
                    )

                with gr.Tab("📈 Analysis"):

                    report = gr.Textbox(
                        label="AI Report",
                        lines=18,
                        interactive=False
                    )

                with gr.Tab("📄 Download"):

                    report_file = gr.File(
                        label="Download Generated Report"
                    )

    gr.Markdown("---")

    with gr.Accordion("⚙ Technical Details", open=False):

        gr.Markdown(
            """
### Model

Communicative Deep Reinforcement Learning

### Framework

PyTorch

### Input

NIfTI / DICOM

### Output

3D Anatomical Landmark Coordinates

### Institution

Research Prototype
"""
        )

    gr.Markdown(
        """
<div class='footer'>
Communicative Reinforcement Learning Landmark Detection • Powered by Gradio
</div>
"""
    )

    def analyse(file):

        fig, text, gif_path = detect(file)

        try:

            image_path = prepare_image(file)

            if image_path.endswith(".nii") or image_path.endswith(".nii.gz"):

                img = nib.load(image_path)

                arr = img.get_fdata()

                info = f"""
Filename : {os.path.basename(image_path)}

Shape : {arr.shape}

Data Type : {arr.dtype}

Dimensions : {len(arr.shape)}

Minimum : {arr.min():.2f}

Maximum : {arr.max():.2f}
"""

            else:

                ds = pydicom.dcmread(image_path)

                info = f"""
Patient : {getattr(ds,'PatientName','Unknown')}

Modality : {getattr(ds,'Modality','Unknown')}

Rows : {ds.Rows}

Columns : {ds.Columns}

Manufacturer :
{getattr(ds,'Manufacturer','Unknown')}
"""

        except:

            info = "Unable to extract metadata."

        report_path = "analysis_report.txt"

        with open(report_path,"w") as f:

            f.write(text)

        return (
            fig,
            text,
            info,
            report_path,
            gif_path
        )

    detect_btn.click(
        analyse,
        inputs=file_input,
        outputs=[
            output_image,
            report,
            metadata,
            report_file,
            path_gif
        ]
    )

demo.launch()