import json
from pathlib import Path

import gradio as gr
import pandas as pd

from phenotype_inference import PhenotypeClassifier
from phenotype_next_tests import PhenotypeNextTestRecommender


BASE_DIR = Path(__file__).parent

MODEL_PATH = BASE_DIR / "phenotype_tinytransformer_v1_temperature_scaled.pt"
REFERENCE_PATH = BASE_DIR / "phenotype_reference_distributions.json"

classifier = PhenotypeClassifier(MODEL_PATH)
recommender = PhenotypeNextTestRecommender(classifier, REFERENCE_PATH)

with open(REFERENCE_PATH, "r", encoding="utf-8") as f:
    reference = json.load(f)

SCHEMA = reference["schema"]
GLOBAL_FIELD_VALUE_COUNTS = reference["global_field_value_counts"]
DEFAULT_RECOMMENDABLE_FIELDS = reference["default_recommendable_fields"]

DISCLAIMER = (
    "This is a phenotype-based genus prediction tool. "
    "It is not a confirmed laboratory identification, diagnostic result, "
    "or replacement for validated microbiology workflows."
)


def get_field_choices(field):
    counter = GLOBAL_FIELD_VALUE_COUNTS.get(field, {})
    values = sorted(counter.keys())

    # Keep dropdowns sensible.
    if field in {"Colony Morphology", "Media Grown On", "Growth Temperature"}:
        return values[:200]

    return values


def clean_value(value):
    if value is None:
        return None

    if isinstance(value, str):
        value = value.strip()
        if value == "":
            return None

    return value


def collect_features(*values):
    features = {}

    for field, value in zip(SCHEMA, values):
        cleaned = clean_value(value)

        if cleaned is None:
            continue

        # Gradio CheckboxGroup may return list.
        if isinstance(cleaned, list):
            cleaned = [str(v).strip() for v in cleaned if str(v).strip()]
            if not cleaned:
                continue
            features[field] = "; ".join(cleaned)
        else:
            features[field] = str(cleaned).strip()

    return features


def confidence_badge(confidence):
    if confidence == "High":
        return "🟢 High"
    if confidence == "Medium":
        return "🟡 Medium"
    return "🔴 Low"


def format_prediction_markdown(prediction):
    lines = []

    lines.append("## Prediction")
    lines.append("")
    lines.append(f"**Top genus:** `{prediction['top_genus']}`")
    lines.append(f"**Probability:** `{prediction['top_probability']:.4f}`")
    lines.append(f"**Margin:** `{prediction['margin']:.4f}`")
    lines.append(f"**Confidence:** {confidence_badge(prediction['confidence'])}")
    lines.append(f"**Distinctness:** `{prediction['distinctness']}`")
    lines.append(f"**Fields used:** `{prediction['num_provided_fields']}`")
    lines.append(f"**Model tokens:** `{prediction['num_model_tokens']}`")
    lines.append(f"**Unknown model tokens:** `{prediction['num_unknown_model_tokens']}`")
    lines.append("")
    lines.append(f"> {DISCLAIMER}")

    return "\n".join(lines)


def ranked_genera_dataframe(prediction):
    rows = []

    for item in prediction["ranked_genera"]:
        rows.append({
            "Rank": item["rank"],
            "Genus": item["genus"],
            "Probability": round(item["probability"], 6),
        })

    return pd.DataFrame(rows)


def recommendations_dataframe(recommendations):
    rows = []

    for i, rec in enumerate(recommendations, start=1):
        likely_outcomes = []

        for value in rec["candidate_values"][:5]:
            likely_outcomes.append(
                f"{value['value']} → {value['top_genus_after']} "
                f"(w={value['estimated_outcome_weight']:.2f}, "
                f"p={value['top_probability_after']:.2f})"
            )

        rows.append({
            "Rank": i,
            "Field": rec["field"],
            "Discriminatory Score": round(rec.get("discriminatory_score", 0), 4),
            "Confirmation Score": round(rec.get("confirmation_score", 0), 4),
            "Model Info Gain": round(rec.get("model_information_gain_bits", 0), 4),
            "Pairwise Separation": round(rec.get("empirical_pairwise_tv_separation", 0), 4),
            "Challenge Rate": round(rec.get("challenge_rate", 0), 4),
            "Evidence Records": rec.get("evidence_records_among_top_genera", 0),
            "Likely Outcomes": " | ".join(likely_outcomes),
        })

    return pd.DataFrame(rows)


def provided_missing_markdown(prediction):
    provided = prediction.get("provided_fields", [])
    missing = prediction.get("missing_fields", [])
    ignored = prediction.get("unknown_input_fields_ignored", [])
    unknown_tokens = prediction.get("unknown_model_tokens", [])

    lines = []
    lines.append("## Input audit")
    lines.append("")
    lines.append("### Provided fields")
    if provided:
        lines.extend([f"- {field}" for field in provided])
    else:
        lines.append("- None")

    lines.append("")
    lines.append("### Missing fields")
    if missing:
        lines.extend([f"- {field}" for field in missing])
    else:
        lines.append("- None")

    if ignored:
        lines.append("")
        lines.append("### Unknown input fields ignored")
        lines.extend([f"- {field}" for field in ignored])

    if unknown_tokens:
        lines.append("")
        lines.append("### Unknown model tokens")
        lines.extend([f"- {token}" for token in unknown_tokens])

    return "\n".join(lines)


def predict_and_recommend(
    top_k,
    n_recommendations,
    top_competing_genera,
    include_context_fields,
    excluded_next_test_fields,
    *field_values,
):
    features = collect_features(*field_values)

    if not features:
        empty_df = pd.DataFrame(columns=["Rank", "Genus", "Probability"])
        empty_rec_df = pd.DataFrame()
        return (
            "Please enter at least one phenotype field.",
            empty_df,
            empty_rec_df,
            empty_rec_df,
            "No input provided.",
            "{}",
        )

    prediction = classifier.predict(features, top_k=int(top_k))

    excluded_next_test_fields = set(excluded_next_test_fields or [])

    if include_context_fields:
        candidate_fields = [
            field for field in SCHEMA
            if field not in features and field not in excluded_next_test_fields
        ]
    else:
        candidate_fields = [
            field for field in DEFAULT_RECOMMENDABLE_FIELDS
            if field not in features and field not in excluded_next_test_fields
        ]

    recommendations = recommender.recommend(
        features,
        n_recommendations=int(n_recommendations),
        top_competing_genera=int(top_competing_genera),
        max_candidate_values_per_field=8,
        include_context_fields=bool(include_context_fields),
        fields_to_consider=candidate_fields,
    )

    prediction_md = format_prediction_markdown(prediction)
    ranked_df = ranked_genera_dataframe(prediction)

    discriminatory_df = recommendations_dataframe(
        recommendations["discriminatory_recommendations"]
    )

    confirmation_df = recommendations_dataframe(
        recommendations["confirmation_recommendations"]
    )

    audit_md = provided_missing_markdown(prediction)

    raw_json = json.dumps(
        {
            "input_features": features,
            "prediction": prediction,
            "recommendations": recommendations,
        },
        indent=2,
    )

    return prediction_md, ranked_df, discriminatory_df, confirmation_df, audit_md, raw_json


def example_achromobacter():
    return {
        "Gram Stain": "Negative",
        "Shape": "Rods",
        "Catalase": "Positive",
        "Oxidase": "Positive",
        "Motility": "Positive",
        "Indole": "Negative",
        "Citrate": "Positive",
        "Urease": "Negative",
        "Growth Temperature": "20//37",
        "Media Grown On": "Blood Agar; MacConkey Agar",
    }


def example_staphylococcus():
    return {
        "Gram Stain": "Positive",
        "Shape": "Cocci",
        "Catalase": "Positive",
        "Oxidase": "Negative",
        "Coagulase": "Positive",
        "DNase": "Positive",
        "Glucose Fermentation": "Positive",
        "Mannitol Fermentation": "Positive",
        "NaCl Tolerant (>=6%)": "Positive",
        "Haemolysis": "Positive",
        "Haemolysis Type": "Beta",
    }


def example_salmonella_like():
    return {
        "Gram Stain": "Negative",
        "Shape": "Rods",
        "Catalase": "Positive",
        "Oxidase": "Negative",
        "Glucose Fermentation": "Positive",
        "Lactose Fermentation": "Negative",
        "Sucrose Fermentation": "Negative",
        "H2S": "Positive",
        "Urease": "Negative",
        "Indole": "Negative",
        "Citrate": "Positive",
        "Motility": "Positive",
        "TSI Pattern": "K/A",
        "Gas Production": "Positive",
    }


EXAMPLES = [
    example_achromobacter(),
    example_staphylococcus(),
    example_salmonella_like(),
]


def values_from_example(example):
    return [example.get(field, None) for field in SCHEMA]


def clear_all():
    return [None for _ in SCHEMA]


with gr.Blocks(title="PhenotypeClassifier", theme=gr.themes.Soft()) as demo:
    gr.Markdown(
        """
# PhenotypeClassifier

Phenotype-based bacterial genus prediction using a calibrated TinyTransformer model.

This demo returns a ranked genus prediction and recommends next tests in two ways:

- **Discriminatory tests**: best for separating the current top competing genera.
- **Confirmation tests**: best for strengthening/checking the current top prediction.

> This is not a confirmed laboratory identification and should not replace validated microbiology workflows.
        """
    )

    with gr.Row():
        with gr.Column(scale=1):
            gr.Markdown("## Input phenotype")

            field_components = []

            with gr.Accordion("Core morphology and basic tests", open=True):
                core_fields = [
                    "Gram Stain",
                    "Shape",
                    "Catalase",
                    "Oxidase",
                    "Colony Morphology",
                    "Haemolysis",
                    "Haemolysis Type",
                    "Growth Temperature",
                    "Media Grown On",
                    "Motility",
                    "Oxygen Requirement",
                ]

                for field in core_fields:
                    if field in {"Colony Morphology", "Media Grown On"}:
                        comp = gr.Textbox(
                            label=field,
                            placeholder="Use semicolons for multiple values, e.g. Smooth; Grey",
                        )
                    elif field == "Growth Temperature":
                        comp = gr.Textbox(
                            label=field,
                            placeholder="e.g. 20//37",
                        )
                    else:
                        comp = gr.Dropdown(
                            label=field,
                            choices=get_field_choices(field),
                            value=None,
                            allow_custom_value=True,
                        )

                    field_components.append((field, comp))

            with gr.Accordion("Biochemical tests", open=True):
                biochemical_fields = [
                    "Indole",
                    "Methyl Red",
                    "VP",
                    "Citrate",
                    "Urease",
                    "H2S",
                    "Nitrate Reduction",
                    "Lysine Decarboxylase",
                    "Ornithine Decarboxylase",
                    "Arginine dihydrolase",
                    "Gelatin Hydrolysis",
                    "Esculin Hydrolysis",
                    "DNase",
                    "ONPG",
                    "Lipase Test",
                    "Coagulase",
                    "TSI Pattern",
                    "Gas Production",
                ]

                for field in biochemical_fields:
                    comp = gr.Dropdown(
                        label=field,
                        choices=get_field_choices(field),
                        value=None,
                        allow_custom_value=True,
                    )
                    field_components.append((field, comp))

            with gr.Accordion("Fermentation and tolerance tests", open=False):
                fermentation_fields = [
                    "Lactose Fermentation",
                    "Glucose Fermentation",
                    "Sucrose Fermentation",
                    "Xylose Fermentation",
                    "Rhamnose Fermentation",
                    "Mannitol Fermentation",
                    "Sorbitol Fermentation",
                    "Maltose Fermentation",
                    "Arabinose Fermentation",
                    "Raffinose Fermentation",
                    "Inositol Fermentation",
                    "Trehalose Fermentation",
                    "NaCl Tolerant (>=6%)",
                ]

                for field in fermentation_fields:
                    comp = gr.Dropdown(
                        label=field,
                        choices=get_field_choices(field),
                        value=None,
                        allow_custom_value=True,
                    )
                    field_components.append((field, comp))

            with gr.Accordion("Other structural features", open=False):
                other_fields = [
                    "Motility Type",
                    "Capsule",
                    "Spore Formation",
                ]

                for field in other_fields:
                    comp = gr.Dropdown(
                        label=field,
                        choices=get_field_choices(field),
                        value=None,
                        allow_custom_value=True,
                    )
                    field_components.append((field, comp))

            # Reorder components to match SCHEMA exactly.
            component_by_field = {field: comp for field, comp in field_components}
            ordered_components = [component_by_field[field] for field in SCHEMA]

            with gr.Row():
                predict_button = gr.Button("Predict and recommend next tests", variant="primary")
                clear_button = gr.Button("Clear")

            with gr.Accordion("Examples", open=False):
                example_buttons = []
                example_buttons.append(gr.Button("Load Achromobacter-like example"))
                example_buttons.append(gr.Button("Load Staphylococcus-like example"))
                example_buttons.append(gr.Button("Load Salmonella-like example"))

        with gr.Column(scale=1):
            gr.Markdown("## Settings")

            top_k = gr.Slider(
                label="Number of ranked genera to show",
                minimum=5,
                maximum=30,
                value=10,
                step=1,
            )

            n_recommendations = gr.Slider(
                label="Number of next-test recommendations",
                minimum=3,
                maximum=10,
                value=5,
                step=1,
            )

            top_competing_genera = gr.Slider(
                label="Number of top competing genera used for next-test simulation",
                minimum=3,
                maximum=10,
                value=5,
                step=1,
            )

            include_context_fields = gr.Checkbox(
                label="Allow context fields as next-test recommendations",
                value=False,
                info="If enabled, fields like Media Grown On, Colony Morphology, and Growth Temperature may be recommended.",
            )

            excluded_next_test_fields = gr.CheckboxGroup(
                label="Exclude these fields from next-test recommendations",
                choices=SCHEMA,
                value=["Oxygen Requirement"],
                info="Useful for hiding fields that are not practical as follow-up tests in your workflow.",
            )

            prediction_md = gr.Markdown()
            ranked_df = gr.Dataframe(label="Ranked genera", interactive=False)
            discriminatory_df = gr.Dataframe(label="Discriminatory next tests", interactive=False)
            confirmation_df = gr.Dataframe(label="Confirmation next tests", interactive=False)

            with gr.Accordion("Input audit", open=False):
                audit_md = gr.Markdown()

            with gr.Accordion("Raw JSON output", open=False):
                raw_json = gr.Code(language="json")

    predict_inputs = [
        top_k,
        n_recommendations,
        top_competing_genera,
        include_context_fields,
        excluded_next_test_fields,
    ] + ordered_components

    predict_outputs = [
        prediction_md,
        ranked_df,
        discriminatory_df,
        confirmation_df,
        audit_md,
        raw_json,
    ]

    predict_button.click(
        fn=predict_and_recommend,
        inputs=predict_inputs,
        outputs=predict_outputs,
    )

    clear_button.click(
        fn=clear_all,
        inputs=[],
        outputs=ordered_components,
    )

    example_buttons[0].click(
        fn=lambda: values_from_example(EXAMPLES[0]),
        inputs=[],
        outputs=ordered_components,
    )

    example_buttons[1].click(
        fn=lambda: values_from_example(EXAMPLES[1]),
        inputs=[],
        outputs=ordered_components,
    )

    example_buttons[2].click(
        fn=lambda: values_from_example(EXAMPLES[2]),
        inputs=[],
        outputs=ordered_components,
    )


if __name__ == "__main__":
    demo.launch()