"""Phase 3: Train & Export -- fine-tune and export ONNX."""

from __future__ import annotations

import argparse
import logging
import sys
import threading
from pathlib import Path

import numpy as np
import streamlit as st
from PIL import Image

from pyzm.train.app import MIN_IMAGES_PER_CLASS, _section_header, _step_expander
from pyzm.train.dataset import YOLODataset
from pyzm.train.trainer import (
    HardwareInfo,
    TrainProgress,
    TrainResult,
    YOLOTrainer,
    adaptive_finetune_params,
)
from pyzm.train.verification import VerificationStore

logger = logging.getLogger("pyzm.train")


# ===================================================================
# PHASE 3: Train & Export
# ===================================================================

def _phase_train(ds: YOLODataset, store: VerificationStore, args: argparse.Namespace) -> None:
    _section_header("&#x1F3AF;", "Train & Export")
    pdir = st.session_state.get("workspace_dir")
    base_model = st.session_state.get("base_model", "yolo11s")
    if not pdir:
        return

    classes = store.build_class_list()
    if not classes:
        # Fall back to classes already saved in the dataset (e.g. from a
        # previous training run or from the YOLO dataset import).
        classes = list(ds.classes) if ds.classes else []
    if not classes:
        reviewed = store.reviewed_images_count()
        if reviewed > 0:
            st.warning(
                "All detections were deleted — no object classes to train on. "
                "Background-only images need at least **some** images with "
                "approved detections to be useful.\n\n"
                "Go back to **Review Detections** and either approve correct "
                "detections or draw boxes on objects the model should learn."
            )
        else:
            st.info("No verified detections yet. Go to **Review Detections** first.")
        if st.button("Go to Review Detections", key="train_go_review"):
            st.session_state["active_phase"] = "review"
            st.rerun()
        return

    # Readiness check — only corrected classes need the image threshold
    needs = store.classes_needing_upload(min_images=MIN_IMAGES_PER_CLASS)
    if needs:
        names = ", ".join(
            f"**{e['class_name']}** ({e['current_images']}/{e['target_images']})"
            for e in needs
        )
        st.warning(f"Need more images for: {names}")
        if st.button("Import More Images", key="train_go_select"):
            st.session_state["active_phase"] = "select"
            st.rerun()

    images = ds.staged_images()
    if len(images) < 2:
        st.info(
            f"You have {len(images)} image(s). We recommend at least "
            f"10\u201320 images per object type for effective training."
        )
        return

    # Shared mutable dict — background thread mutates contents,
    # main thread reads.  Lives in session_state so it survives reruns.
    shared: dict = st.session_state.get("_train_shared", {})

    # Restore previous training results from disk if no active session
    if not shared.get("active") and shared.get("result") is None:
        disk_result = TrainResult.from_disk(Path(pdir))
        if disk_result is not None:
            shared["result"] = disk_result
            # Check for existing ONNX export
            best_onnx = Path(pdir) / "runs" / "train" / "weights" / "best.onnx"
            if best_onnx.exists():
                shared["onnx_path"] = str(best_onnx)
            st.session_state["_train_shared"] = shared
            if "classes" not in st.session_state:
                st.session_state["classes"] = classes

    result: TrainResult | None = shared.get("result")
    training_active = shared.get("active", False)
    train_done = result is not None

    # === Step 1: Configure training ===
    with _step_expander("Configure training", done=train_done or training_active):
        hw = YOLOTrainer.detect_hardware()

        col1, col2, col3, col4 = st.columns(4)
        with col1:
            epochs = st.number_input(
                "Epochs", min_value=1, max_value=300, value=50, step=10,
                help="How many times to review all images during training. 50 is usually good for fine-tuning.",
            )
        with col2:
            batch = st.number_input(
                "Batch", min_value=1, max_value=128, value=hw.suggested_batch,
                help="Images processed at once. Auto-detected for your hardware.",
            )
        with col3:
            imgsz = st.selectbox(
                "Image size", [416, 640], index=1,
                help="Resolution for training. 640 is standard.",
            )
        with col4:
            st.caption(hw.display)

        # Determine new vs known classes for smart defaults
        model_class_names: list[str] = (
            st.session_state.get("model_class_names")
            or ds.get_setting("model_class_names")
            or []
        )
        new_classes = [c for c in classes if c not in model_class_names]
        known_classes = [c for c in classes if c in model_class_names]
        has_new = bool(new_classes)

        # Restore previous selection; fall back to heuristic
        _saved_ft = ds.get_setting("ft_mode", "")
        if _saved_ft in ("new_class", "refine"):
            _ft_default_idx = 0 if _saved_ft == "new_class" else 1
        else:
            _ft_default_idx = 0 if has_new else 1

        # Contextual recommendation
        n_images = len(images)
        if has_new:
            new_list = ", ".join(f"**{c}**" for c in new_classes)
            st.info(
                f":material/lightbulb: **Recommended: New class** — your dataset "
                f"includes classes the base model doesn't know ({new_list}). "
                f"New class mode uses minimal augmentation so the model learns "
                f"clean representations first."
            )
        elif n_images < 50:
            st.info(
                f":material/lightbulb: **Recommended: Refine existing** — all "
                f"your classes are already known to the base model. With only "
                f"**{n_images}** images, improvements may be modest — consider "
                f"adding more images for better results."
            )
        else:
            st.info(
                f":material/lightbulb: **Recommended: Refine existing** — all "
                f"your classes are already known to the base model. You're "
                f"teaching it to detect them better in your specific environment."
            )

        ft_mode = st.radio(
            "Fine-tuning mode",
            options=["new_class", "refine"],
            index=_ft_default_idx,
            format_func=lambda x: "New class" if x == "new_class" else "Refine existing",
            horizontal=True,
        )

        adaptive = adaptive_finetune_params(len(images), mode=ft_mode)
        st.caption(
            f":material/tune: Adaptive fine-tuning ({adaptive['tier']} dataset, "
            f"{adaptive['mode']} mode, {len(images)} images): "
            f"freeze={adaptive['freeze']}, lr0={adaptive['lr0']}, "
            f"mosaic={adaptive['mosaic']}, erasing={adaptive['erasing']}, "
            f"patience={adaptive['patience']}"
        )

        if not training_active:
            all_ready = not needs
            deploy_tip = ""
            if ft_mode == "new_class":
                deploy_tip = (
                    " **Deployment**: add this model alongside your base model "
                    "in `objectconfig.yml` and use "
                    "`same_model_sequence_strategy: union` to combine detections."
                )
            st.info(
                f"This fine-tuned model will **only** detect the following "
                f"{len(classes)} object(s): **{', '.join(classes)}**. "
                f"It will not retain the base model's original classes. "
                f"To detect other objects, run the base model alongside this one."
                + deploy_tip
            )
            if hw.device == "cpu":
                st.warning(
                    "Training on CPU will be significantly slower than GPU. "
                    "Consider using a machine with a CUDA-capable GPU for faster training."
                )
            hw_factor = 1 if hw.device != "cpu" else 8
            est_minutes = int((len(images) * epochs) / (batch * 60) * hw_factor)
            if est_minutes > 0:
                st.caption(f":material/timer: Estimated training time: ~{est_minutes} minutes (rough estimate)")
            if st.button(":material/rocket_launch: Start Training", type="primary", disabled=not all_ready):
                class_name_to_id = {c: i for i, c in enumerate(classes)}
                ds.set_classes(classes)
                ds.set_setting("ft_mode", ft_mode)
                yaml_path = ds.generate_yaml()

                trainer = YOLOTrainer(
                    base_model=base_model,
                    project_dir=Path(pdir),
                    device=hw.device,
                )
                shared = {
                    "active": True,
                    "progress": TrainProgress(
                        total_epochs=epochs, message="Preparing dataset...",
                    ),
                    "result": None,
                    "log": [],
                    "adaptive": adaptive,
                }
                import time as _time
                shared["start_time"] = _time.time()
                shared["image_count"] = len(images)
                st.session_state["_train_shared"] = shared
                st.session_state["trainer"] = trainer
                st.session_state["classes"] = classes

                def _run(_s: dict = shared) -> None:
                    import_classes = ds.get_setting("import_classes")
                    need_rewrite = (
                        store.has_modifications()
                        or import_classes != classes
                    )

                    n = len(images)
                    if need_rewrite:
                        for i, img_path in enumerate(images):
                            anns = store.finalized_annotations(
                                img_path.name, class_name_to_id,
                            )
                            ds.update_annotations(img_path.name, anns)
                            if n > 100 and (i + 1) % 50 == 0:
                                _s["progress"] = TrainProgress(
                                    total_epochs=epochs,
                                    message=f"Writing annotations {i + 1}/{n}",
                                )

                    split_map = ds.get_setting("split_map") or {}
                    has_split = ds._train_images.exists() and any(ds._train_images.iterdir())
                    if need_rewrite or not has_split:
                        _s["progress"] = TrainProgress(
                            total_epochs=epochs,
                            message="Splitting into train/val...",
                        )
                        ds.split(_s["adaptive"]["val_ratio"])

                    _s["progress"] = TrainProgress(
                        total_epochs=epochs, message="Loading model...",
                    )

                    def _cb(p: TrainProgress) -> None:
                        _s["progress"] = p

                    class _TrainLogHandler(logging.Handler):
                        def emit(self, record: logging.LogRecord) -> None:
                            log = _s["log"]
                            log.append(self.format(record))
                            if len(log) > 200:
                                del log[:-200]

                    class _StreamCapture:
                        def __init__(self, original):
                            self._original = original
                        def write(self, text):
                            self._original.write(text)
                            if text.strip():
                                log = _s["log"]
                                log.append(text.rstrip())
                                if len(log) > 200:
                                    del log[:-200]
                        def flush(self):
                            self._original.flush()
                        def __getattr__(self, name):
                            return getattr(self._original, name)

                    handler = _TrainLogHandler()
                    handler.setFormatter(logging.Formatter("%(message)s"))
                    ul_logger = logging.getLogger("ultralytics")
                    ul_logger.addHandler(handler)
                    ul_logger.setLevel(logging.INFO)
                    old_stdout, old_stderr = sys.stdout, sys.stderr
                    sys.stdout = _StreamCapture(old_stdout)
                    sys.stderr = _StreamCapture(old_stderr)
                    try:
                        _adapt = _s["adaptive"]
                        _train_extra = {
                            k: _adapt[k] for k in (
                                "freeze", "lr0", "patience", "cos_lr",
                                "mosaic", "erasing", "scale", "mixup", "close_mosaic",
                            )
                        }
                        r = trainer.train(
                            dataset_yaml=yaml_path, epochs=epochs,
                            batch=batch, imgsz=imgsz, progress_callback=_cb,
                            **_train_extra,
                        )
                        # Auto-export ONNX
                        if r.best_model:
                            try:
                                onnx_path = trainer.export_onnx()
                                _s["onnx_path"] = str(onnx_path)
                            except Exception as exc:
                                _s["onnx_error"] = str(exc)
                        _s["result"] = r
                    except Exception as exc:
                        _s["progress"] = TrainProgress(finished=True, error=str(exc))
                    finally:
                        sys.stdout, sys.stderr = old_stdout, old_stderr
                        ul_logger.removeHandler(handler)
                        _s["active"] = False

                threading.Thread(target=_run, daemon=True).start()
                st.rerun()

    # === Training progress (not in expander — must stay visible) ===
    if training_active:
        st.warning("Do not close this browser tab while training is in progress. Training will be lost.")
        p: TrainProgress = shared.get("progress") or TrainProgress()
        if p.total_epochs > 0:
            pct = max(0.0, min(1.0, p.epoch / p.total_epochs))
            st.progress(pct, text=p.message or f"Epoch {p.epoch}/{p.total_epochs}")

        import time as _time
        start = shared.get("start_time")
        if start:
            elapsed = _time.time() - start
            mins, secs = divmod(int(elapsed), 60)
            hrs, mins = divmod(mins, 60)
            elapsed_str = f"{hrs}:{mins:02d}:{secs:02d}" if hrs else f"{mins}:{secs:02d}"
        else:
            elapsed_str = "0:00"
        img_count = shared.get("image_count", 0)

        t1, t2 = st.columns(2)
        t1.markdown(f":material/timer: **Elapsed Time:** {elapsed_str}")
        t2.markdown(f":material/image: **Training Images:** {img_count}")

        m1, m2, m3 = st.columns(3)
        m1.metric("Epoch", f"{p.epoch}/{p.total_epochs}")

        prev_box = st.session_state.get("_prev_box_loss")
        if prev_box is not None and p.box_loss > 0:
            if p.box_loss < prev_box - 0.001:
                box_label = "Box accuracy: Improving"
                box_color = "green"
            elif p.box_loss > prev_box + 0.001:
                box_label = "Box accuracy: Declining"
                box_color = "red"
            else:
                box_label = "Box accuracy: Stable"
                box_color = "orange"
        else:
            box_label = "Box accuracy: Starting"
            box_color = "gray"
        if p.epoch > 0:
            st.session_state["_prev_box_loss"] = p.box_loss
        m2.markdown(f":{box_color}[{box_label}]")

        mAP_pct = p.mAP50 * 100
        if mAP_pct >= 80:
            quality = "Excellent"
        elif mAP_pct >= 60:
            quality = "Good"
        elif mAP_pct >= 30:
            quality = "Fair"
        else:
            quality = "Poor"
        m3.metric("Detection quality", f"{mAP_pct:.0f}% ({quality})")

        with st.expander("Technical Details"):
            td1, td2, td3, td4 = st.columns(4)
            td1.metric("Epoch", f"{p.epoch}/{p.total_epochs}")
            td2.metric("Box Loss", f"{p.box_loss:.4f}")
            td3.metric("Cls Loss", f"{p.cls_loss:.4f}")
            td4.metric("mAP50", f"{p.mAP50:.3f}")

        if st.button(":material/stop_circle: Stop Training"):
            trainer = st.session_state.get("trainer")
            if trainer:
                trainer.request_stop()

        train_log = shared.get("log", [])
        if train_log:
            st.code("\n".join(train_log[-30:]), language=None)

        if p.error:
            st.error(p.error)
        # Keep polling while the training thread is alive — p.finished may
        # be True (trainer returned) while the thread is still exporting
        # ONNX and writing results.  Stop only once active becomes False.
        if shared.get("active"):
            if p.finished:
                st.info("Training complete — exporting model...")
            import time
            time.sleep(2)
            st.rerun()

    # === Step 2: Training results ===
    if train_done:
        # Prominent completion banner (always visible)
        onnx_path = shared.get("onnx_path")
        onnx_error = shared.get("onnx_error")
        mAP_pct = result.final_mAP50 * 100
        early_stop = (
            result.best_epoch > 0
            and result.best_epoch < result.total_epochs
        )
        early_note = (
            f" Training stopped early at epoch {result.best_epoch}/{result.total_epochs} "
            f"(no further improvement detected)."
            if early_stop else ""
        )
        st.success(
            f":material/check_circle: **Training complete!** "
            f"mAP50: **{mAP_pct:.0f}%** — "
            f"Model size: **{result.model_size_mb:.1f} MB** — "
            f"Time: **{result.elapsed_seconds / 60:.1f} min**"
            + early_note
        )
        if onnx_path:
            st.info(f":material/download_done: ONNX model ready: `{onnx_path}`")
        elif onnx_error:
            st.warning(f"ONNX export failed: {onnx_error}")

        with _step_expander(
            "Training results",
            done=True,
            detail=f"mAP50: {result.final_mAP50:.3f}",
        ):
            r1, r2, r3, r4 = st.columns(4)
            r1.metric("mAP50", f"{result.final_mAP50:.3f}")
            r2.metric("mAP50-95", f"{result.final_mAP50_95:.3f}")
            r3.metric("Size", f"{result.model_size_mb:.1f} MB")
            r4.metric("Time", f"{result.elapsed_seconds / 60:.1f} min")

            best_ep_label = (
                f"Epoch {result.best_epoch}/{result.total_epochs}"
                if result.best_epoch > 0
                else f"{result.total_epochs} epochs"
            )
            st.caption(f":material/star: Best model from: **{best_ep_label}**")

            if result.per_class:
                import pandas as pd

                rows = []
                for cls_name, cm in sorted(result.per_class.items()):
                    rows.append({
                        "Class": cls_name,
                        "Precision": f"{cm.precision:.3f}",
                        "Recall": f"{cm.recall:.3f}",
                        "AP@50": f"{cm.ap50:.3f}",
                        "AP@50-95": f"{cm.ap50_95:.3f}",
                    })
                st.dataframe(pd.DataFrame(rows), hide_index=True, width="stretch")

            pdir = st.session_state.get("workspace_dir")
            if pdir:
                st.caption(f"Dataset: `{pdir}`")
                _training_analysis(result, Path(pdir) / "runs" / "train")

            _test_image_panel(args)


def _training_analysis(result: TrainResult, train_dir: Path) -> None:
    """Show interpretive guidance and diagnostic images after training."""
    _section_header("&#x1F4CA;", "Training Analysis")

    # -- (a) Interpretive guidance --
    mAP = result.final_mAP50
    if mAP >= 0.8:
        st.success(f"**Excellent** — mAP50 {mAP:.2f} indicates strong detection quality.")
    elif mAP >= 0.6:
        st.info(f"**Good** — mAP50 {mAP:.2f}. Model performs well; more data may push it higher.")
    elif mAP >= 0.3:
        st.warning(f"**Moderate** — mAP50 {mAP:.2f}. Consider adding more diverse training images.")
    else:
        st.warning(f"**Poor** — mAP50 {mAP:.2f}. The model needs significantly more training data.")

    # Per-class weak spots
    weak = [name for name, cm in result.per_class.items() if cm.ap50 < 0.5]
    if weak:
        st.info(
            f"**Weak classes** (AP50 < 0.5): {', '.join(sorted(weak))}. "
            "Consider adding more training images for these."
        )

    # Overfitting hint
    if (
        result.best_epoch > 0
        and result.total_epochs > 0
        and result.best_epoch < result.total_epochs * 0.5
    ):
        st.info(
            f"Best model was at epoch {result.best_epoch}/{result.total_epochs} "
            "(early in training). The model may have overfit — "
            "try fewer epochs or more training data."
        )

    # -- (b) Training curves --
    results_png = train_dir / "results.png"
    if results_png.exists():
        st.markdown("##### Training Curves")
        with st.expander("How to read these curves"):
            st.markdown(
                "**Loss curves** (box_loss, cls_loss, dfl_loss): should decrease "
                "and flatten. If training loss keeps dropping but validation loss "
                "rises, the model is overfitting — stop earlier or add more data.\n\n"
                "**Precision & Recall**: both should rise and stabilize near 1.0. "
                "Low precision = too many false detections; low recall = missing "
                "real objects.\n\n"
                "**mAP50 / mAP50-95**: the main quality scores — higher is better. "
                "A plateau means more epochs won't help; more diverse data will."
            )
        st.image(str(results_png), width="stretch")

    # -- (c) Confusion matrix --
    cm_norm = train_dir / "confusion_matrix_normalized.png"
    cm_plain = train_dir / "confusion_matrix.png"
    cm_path = cm_norm if cm_norm.exists() else cm_plain if cm_plain.exists() else None
    if cm_path:
        st.markdown("##### Which Objects Get Mixed Up")
        with st.expander("How to read the confusion matrix"):
            st.markdown(
                "Each row is a **true** class, each column is a **predicted** class. "
                "Bright diagonal = correct predictions. Off-diagonal cells show "
                "which classes get confused with each other.\n\n"
                "A **background** row/column means missed detections (false negatives) "
                "or phantom detections (false positives). If a class has a high "
                "background score, the model needs more examples of that class."
            )
        st.image(str(cm_path), width="stretch")

    # -- (d) Evaluation curves --
    f1_path = train_dir / "F1_curve.png"
    pr_path = train_dir / "PR_curve.png"
    if f1_path.exists() or pr_path.exists():
        with st.expander("Detection Balance"):
            st.markdown(
                "**F1 Curve**: shows the balance between precision and recall at "
                "each confidence threshold. The peak is the optimal threshold — "
                "a sharp, high peak (close to 1.0) is ideal.\n\n"
                "**PR Curve**: precision vs. recall trade-off. A curve that hugs "
                "the top-right corner is a strong model. Area under the curve "
                "(AUC) equals AP — higher is better."
            )
            c1, c2 = st.columns(2)
            if f1_path.exists():
                c1.image(str(f1_path), caption="F1 Curve", width="stretch")
            if pr_path.exists():
                c2.image(str(pr_path), caption="PR Curve", width="stretch")

    # -- (e) Validation samples --
    val_labels = train_dir / "val_batch0_labels.jpg"
    val_preds = train_dir / "val_batch0_pred.jpg"
    if val_labels.exists() or val_preds.exists():
        with st.expander("Validation Samples"):
            st.markdown(
                "**Ground Truth** shows the actual labels. **Predictions** shows "
                "what the model detected. Compare them — missed objects or wrong "
                "labels indicate classes that need more training data."
            )

            @st.dialog("Validation Sample", width="large")
            def _show_full(path: str, caption: str):
                st.image(path, caption=caption, width="stretch")

            c1, c2 = st.columns(2)
            if val_labels.exists():
                c1.image(str(val_labels), caption="Ground Truth", width="stretch")
                if c1.button("Expand", key="expand_gt", icon=":material/zoom_in:"):
                    _show_full(str(val_labels), "Ground Truth")
            if val_preds.exists():
                c2.image(str(val_preds), caption="Predictions", width="stretch")
                if c2.button("Expand", key="expand_pred", icon=":material/zoom_in:"):
                    _show_full(str(val_preds), "Predictions")


def _test_image_panel(args: argparse.Namespace) -> None:
    """Let the user test the trained model on an image."""
    pdir = st.session_state.get("workspace_dir")
    base_model = st.session_state.get("base_model", "yolo11s")
    if not pdir:
        return

    weights_dir = Path(pdir) / "runs" / "train" / "weights"
    if not (weights_dir / "best.pt").exists() and not (weights_dir / "best.onnx").exists():
        return

    classes = st.session_state.get("classes", [])
    st.info(
        f"This fine-tuned model only detects: **{', '.join(classes)}**. "
        f"To also detect standard objects (person, car, etc.), add both models "
        f"to your `objectconfig.yml` and use `same_model_sequence_strategy: union` "
        f"to merge results."
    )

    min_conf = st.slider(
        "Confidence threshold",
        min_value=0.01, max_value=1.0, value=0.3, step=0.05,
        help="Detections below this confidence are discarded.",
        key="test_min_confidence",
    )

    test_file = st.file_uploader("Test image", type=["jpg", "jpeg", "png"], key="test_img")
    if test_file:
        trainer = YOLOTrainer(base_model=base_model, project_dir=Path(pdir))
        pil_img = Image.open(test_file).convert("RGB")
        img_array = np.array(pil_img)
        try:
            dets = trainer.evaluate(
                img_array[..., ::-1],
                processor=args.processor,
                min_confidence=min_conf,
            )
            from PIL import ImageDraw, ImageFont

            draw_img = pil_img.copy()
            draw = ImageDraw.Draw(draw_img)

            # Scale font size to ~2% of image height (min 14, max 40)
            font_size = max(14, min(40, draw_img.height // 50))
            try:
                font = ImageFont.truetype("/usr/share/fonts/truetype/dejavu/DejaVuSans-Bold.ttf", font_size)
            except (OSError, IOError):
                font = ImageFont.load_default(size=font_size)

            box_width = max(2, draw_img.height // 300)
            for d in dets:
                x1, y1, x2, y2 = d["bbox"]
                label = f"{d['label']} {d['confidence']:.0%}"
                draw.rectangle([x1, y1, x2, y2], outline="lime", width=box_width)
                # Draw text with dark background for readability
                bbox = font.getbbox(label)
                tw, th = bbox[2] - bbox[0], bbox[3] - bbox[1]
                ty = max(0, y1 - th - 4)
                draw.rectangle([x1, ty, x1 + tw + 4, ty + th + 4], fill=(0, 0, 0, 180))
                draw.text((x1 + 2, ty + 2), label, fill="lime", font=font)
            st.image(draw_img, width="stretch")
            if dets:
                st.caption(", ".join(f"{d['label']} {d['confidence']:.0%}" for d in dets))
            else:
                st.caption("No detections")
        except Exception as exc:
            st.error(str(exc))
