import argparse
from pathlib import Path

import gradio as gr
import numpy as np
import torch
import torch.nn.functional as F

from augmentations import get_val_transforms
from model import DeepSeeNet


N_CLASSES = {
    "ADVAMD": 2,
    "DRUS": 3,
    "PIG": 2,
}

LABELS = {
    "ADVAMD": ["no_late_amd", "late_amd"],
    "DRUS": ["small_none", "medium", "large"],
    "PIG": ["no_pigment", "pigment"],
}


APP_CSS = """
#terms-overlay {
    position: fixed;
    inset: 0;
    z-index: 9999;
    display: flex;
    align-items: center;
    justify-content: center;
    padding: 32px;
    background: rgba(255, 255, 255, 0.72);
    backdrop-filter: blur(3px);
    -webkit-backdrop-filter: blur(3px);
}

#terms-card {
    width: min(720px, calc(100vw - 64px));
    max-height: 86vh;
    overflow-y: auto;
    padding: 24px 28px;
    border: 1px solid #d9d9df;
    border-radius: 8px;
    background: #f3f3f5;
    box-shadow: 0 12px 36px rgba(0, 0, 0, 0.12);
    color: #1f2933;
}

#terms-card h1 {
    margin-top: 0;
    font-size: 1.55rem;
}

#terms-card h2 {
    margin-top: 1.35rem;
    font-size: 1.15rem;
}

#terms-card p,
#terms-card ul {
    font-size: 0.95rem;
    line-height: 1.55;
    margin-left: 1.1rem;
}

#terms-card li {
    margin-bottom: 0.35rem;
}

#acknowledge-btn {
    width: 100%;
    margin-top: 16px;
    padding: 10px 14px;
    border: none;
    border-radius: 6px;
    background: #2563eb;
    color: white;
    font-weight: 600;
    cursor: pointer;
}

#acknowledge-btn:hover {
    background: #1d4ed8;
}
"""

TERMS_AND_CONDITIONS_MD = """
# DeepSeeNet Demo — Research Use Only

This demo is a modernized research implementation based on the original DeepSeeNet work. It is provided solely for research, educational, and reproducibility purposes.

## Not for Clinical Use

This Gradio demo and its associated model outputs are provided for research, educational, and reproducibility purposes only, and are not intended for clinical diagnosis, treatment, screening, triage, patient management, or any other medical decision-making purpose.

This demo implementation has not been clinically validated for routine patient care and has not been reviewed, cleared, or approved by any regulatory authority, including the U.S. Food and Drug Administration.

## Data Privacy

Do not upload protected health information, personally identifiable information, or patient data unless you have the necessary rights, permissions, and safeguards to do so.

## User Responsibility

By using this demo, you acknowledge that:

- You understand this demo is for research and educational purposes only.
- You will not use the outputs for diagnosis, treatment, screening, triage, or clinical decision-making.
- You are responsible for how you interpret and use any results generated by the demo.
- You will not upload protected health information or identifiable patient data unless you are authorized to do so.

## Attribution

This project is a modernized reimplementation of DeepSeeNet. Please cite the original DeepSeeNet publication when using this project:

Peng Y, Dharssi S, Chen Q, Keenan TD, Agrón E, Wong WT, Chew EY, Lu Z. DeepSeeNet: a deep learning model for automated classification of patient-based age-related macular degeneration severity from color fundus photographs. Ophthalmology. 2019 Apr 1;126(4):565-75.

## License and Warranty Disclaimer

Permission is hereby granted, free of charge, to any person obtaining a copy of this software and associated documentation files (the “Software”), to deal in the Software without restriction, including without limitation the rights to use, copy, modify, merge, publish, distribute, sublicense, and/or sell copies of the Software, and to permit persons to whom the Software is furnished to do so, subject to the following conditions:

The above copyright notice and this permission notice shall be included in all copies or substantial portions of the Software.

THE SOFTWARE IS PROVIDED “AS IS”, WITHOUT WARRANTY OF ANY KIND, EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE SOFTWARE.
"""

ACKNOWLEDGEMENT_TEXT = "I acknowledge and continue"

TERMS_OVERLAY_HTML = f"""
<div id="terms-overlay">
  <div id="terms-card">
    <h1>DeepSeeNet Demo — Research Use Only</h1>

    <p>This demo is a modernized research implementation based on the original DeepSeeNet work. It is provided solely for research, educational, and reproducibility purposes.</p>

    <h2>Not for Clinical Use</h2>
    <p>This Gradio demo and its associated model outputs are provided for research, educational, and reproducibility purposes only, and are not intended for clinical diagnosis, treatment, screening, triage, patient management, or any other medical decision-making purpose.</p>
    <p>This demo implementation has not been clinically validated for routine patient care and has not been reviewed, cleared, or approved by any regulatory authority, including the U.S. Food and Drug Administration.</p>

    <h2>Data Privacy</h2>
    <p>Do not upload protected health information, personally identifiable information, or patient data unless you have the necessary rights, permissions, and safeguards to do so.</p>

    <h2>User Responsibility</h2>
    <p>By using this demo, you acknowledge that:</p>
    <ul>
      <li>You understand this demo is for research and educational purposes only.</li>
      <li>You will not use the outputs for diagnosis, treatment, screening, triage, or clinical decision-making.</li>
      <li>You are responsible for how you interpret and use any results generated by the demo.</li>
      <li>You will not upload protected health information or identifiable patient data unless you are authorized to do so.</li>
    </ul>

    <h2>Attribution</h2>
    <p>Please cite the original DeepSeeNet publication when using this project:</p>
    <p>Peng Y, Dharssi S, Chen Q, Keenan TD, Agrón E, Wong WT, Chew EY, Lu Z. DeepSeeNet: a deep learning model for automated classification of patient-based age-related macular degeneration severity from color fundus photographs. Ophthalmology. 2019 Apr 1;126(4):565-75.</p>

    <h2>License and Warranty Disclaimer</h2>
    <p>Permission is hereby granted, free of charge, to any person obtaining a copy of this software and associated documentation files (the “Software”), to deal in the Software without restriction, including without limitation the rights to use, copy, modify, merge, publish, distribute, sublicense, and/or sell copies of the Software, and to permit persons to whom the Software is furnished to do so, subject to the following conditions:</p>
    <p>The above copyright notice and this permission notice shall be included in all copies or substantial portions of the Software.</p>
    <p>THE SOFTWARE IS PROVIDED “AS IS”, WITHOUT WARRANTY OF ANY KIND, EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE SOFTWARE.</p>

    <button id="acknowledge-btn" onclick="document.getElementById('terms-overlay').remove();">{ACKNOWLEDGEMENT_TEXT}</button>
  </div>
</div>
"""


class AlbumentationsTransform:
    def __init__(self, transform):
        self.transform = transform

    def __call__(self, image):
        return self.transform(image=np.asarray(image))["image"]


def parse_args():
    parser = argparse.ArgumentParser()
    parser.add_argument("--checkpoint-folder", default="./checkpoints")
    parser.add_argument("--backbone", default="inception_v3")
    parser.add_argument("--image-size", type=int, default=1024)
    parser.add_argument("--server-name", default="127.0.0.1")
    parser.add_argument("--server-port", type=int, default=7860)
    parser.add_argument("--share", action="store_true")
    return parser.parse_args()


def load_model(path, task, backbone, device):
    checkpoint = torch.load(path, map_location=device)
    checkpoint_args = checkpoint.get("args", {})

    model = DeepSeeNet(
        n_classes=N_CLASSES[task],
        backbone=checkpoint_args.get("backbone", backbone),
        pretrained=False,
    ).to(device)

    model.load_state_dict(checkpoint["model"])
    model.eval()
    return model


def load_image(image, transform, device):
    if image is None:
        raise ValueError("Please upload both left and right images.")

    image = image.convert("RGB")
    return transform(image).unsqueeze(0).to(device)


@torch.no_grad()
def predict(model, image, task):
    logits = model(image)[0].detach().cpu()
    probs = F.softmax(logits, dim=0)
    pred = int(torch.argmax(logits).item())

    return {
        "prediction": pred,
        "label": LABELS[task][pred],
        "confidence": float(probs[pred]),
        "probabilities": {
            LABELS[task][i]: float(probs[i])
            for i in range(len(LABELS[task]))
        },
    }


def simplified_score(scores):
    if scores["ADVAMD"]["left"]["prediction"] == 1 or scores["ADVAMD"]["right"]["prediction"] == 1:
        return 5

    score = 0
    score += scores["PIG"]["left"]["prediction"] == 1
    score += scores["PIG"]["right"]["prediction"] == 1
    score += scores["DRUS"]["left"]["prediction"] == 2
    score += scores["DRUS"]["right"]["prediction"] == 2
    score += (
        scores["DRUS"]["left"]["prediction"] == 1
        and scores["DRUS"]["right"]["prediction"] == 1
    )

    return int(min(score, 5))


def format_probs(probabilities):
    return " | ".join(
        f"{label}: {prob:.3f}"
        for label, prob in probabilities.items()
    )


def model_info(args, device):
    return (
        "# DeepSeeNet\n\n"
        f"Model: `{args.backbone}`  |  "
        f"Input size: `{args.image_size} × {args.image_size}`  |  "
        f"Device: `{device.type}`"
        # f"Checkpoint folder: `{args.checkpoint_folder}`"
    )


def make_app(args):
    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
    checkpoint_folder = Path(args.checkpoint_folder)
    transform = AlbumentationsTransform(get_val_transforms(args.image_size))

    models = {
        "ADVAMD": load_model(checkpoint_folder / "advamd.pt", "ADVAMD", args.backbone, device),
        "DRUS": load_model(checkpoint_folder / "drus.pt", "DRUS", args.backbone, device),
        "PIG": load_model(checkpoint_folder / "pig.pt", "PIG", args.backbone, device),
    }

    def run(left_image, right_image):
        left = load_image(left_image, transform, device)
        right = load_image(right_image, transform, device)

        scores = {}
        for task, model in models.items():
            scores[task] = {
                "left": predict(model, left, task),
                "right": predict(model, right, task),
            }

        score = simplified_score(scores)

        summary_rows = [
            ["AREDS simplified score", score],
            ["Left eye", f"{scores['DRUS']['left']['label']}, {scores['PIG']['left']['label']}, {scores['ADVAMD']['left']['label']}"],
            ["Right eye", f"{scores['DRUS']['right']['label']}, {scores['PIG']['right']['label']}, {scores['ADVAMD']['right']['label']}"],
        ]

        detail_rows = []
        for task in ["ADVAMD", "DRUS", "PIG"]:
            for eye in ["left", "right"]:
                result = scores[task][eye]
                detail_rows.append(
                    [
                        task,
                        eye,
                        result["label"],
                        f"{result['confidence']:.3f}",
                        format_probs(result["probabilities"]),
                    ]
                )

        return summary_rows, detail_rows

    with gr.Blocks(title="DeepSeeNet", css=APP_CSS) as demo:
        # Main demo is rendered normally from the start. The terms are only an overlay,
        # so once acknowledged the app behaves as if the landing page was never there.
        gr.Markdown(model_info(args, device))

        with gr.Row():
            left_image = gr.Image(type="pil", label="Left image")
            right_image = gr.Image(type="pil", label="Right image")

        button = gr.Button("Run")

        summary = gr.Dataframe(
            headers=["Item", "Result"],
            label="Summary",
        )
        details = gr.Dataframe(
            headers=["Task", "Eye", "Prediction", "Confidence", "Probabilities"],
            label="Model outputs",
        )

        # Pure front-end overlay: this does not participate in the Gradio layout.
        # Clicking the button removes only the overlay, leaving the underlying app unchanged.
        gr.HTML(TERMS_OVERLAY_HTML)

        button.click(
            run,
            inputs=[left_image, right_image],
            outputs=[summary, details],
        )


    return demo


def main():
    args = parse_args()
    demo = make_app(args)
    demo.launch(
        # server_name=args.server_name,
        # server_port=args.server_port,
        # share=args.share,
    )


if __name__ == "__main__":
    main()