# @launcher_title: 1D U-Net Trainer v1.3
# @launcher_group: UNet
# @launcher_order: 11
# @launcher_bg: #0E8A16
# @launcher_fg: #FFFFFF
# @launcher_desc: U-Net学習 v1.3 DirectML/ONNX direct save
# -*- coding: utf-8 -*-

import datetime
import json
import os
import re
import threading
import time
import xml.etree.ElementTree as ET

import numpy as np
import torch
import torch.nn as nn
from torch.utils.data import Dataset, DataLoader

import tkinter as tk
from tkinter import filedialog, messagebox
from tkinter import ttk

from matplotlib.backends.backend_tkagg import FigureCanvasTkAgg
from matplotlib.figure import Figure

try:
    import torch_directml
except Exception as _torch_directml_error:
    torch_directml = None
    TORCH_DIRECTML_IMPORT_ERROR = _torch_directml_error
else:
    TORCH_DIRECTML_IMPORT_ERROR = None


def strip_ns(tag: str) -> str:
    return tag.split("}", 1)[-1]


def get_child_text(elem, name, default=None):
    for child in elem:
        if strip_ns(child.tag) == name:
            txt = (child.text or "").strip()
            return txt if txt != "" else default
    return default


def safe_compound_name(name: str) -> str:
    return "".join(c if c.isalnum() else "_" for c in (name or ""))


def wildcard_regex_from_safe_name(name: str):
    pattern = "^" + "".join(".*" if c == "_" else re.escape(c) for c in (name or "")) + "$"
    return re.compile(pattern, re.IGNORECASE)


def find_compound_metadata(compound_name: str, metadata_by_compound: dict):
    if not metadata_by_compound:
        return {}, ""

    if compound_name in metadata_by_compound:
        return metadata_by_compound[compound_name], compound_name

    safe_target = safe_compound_name(compound_name)
    for key, value in metadata_by_compound.items():
        if safe_compound_name(key).lower() == safe_target.lower():
            return value, key

    target_regex = wildcard_regex_from_safe_name(safe_target)
    for key, value in metadata_by_compound.items():
        safe_key = safe_compound_name(key)
        if target_regex.match(key) or target_regex.match(safe_key):
            return value, key

    for key, value in metadata_by_compound.items():
        key_regex = wildcard_regex_from_safe_name(safe_compound_name(key))
        if key_regex.match(compound_name) or key_regex.match(safe_target):
            return value, key

    return {}, ""


PTH_TEXT_METADATA_FIELDS = [
    "Comment",
    "Comment2",
    "Q1",
    "Q2",
    "Q3",
    "Chromato_color",
    "ISTD",
    "Area_value",
    "Intensity_threshold",
    "Calib",
    "IsChecked",
]


def state_dict_to_cpu(model):
    return {
        key: value.detach().cpu()
        for key, value in model.state_dict().items()
    }


def add_onnx_metadata(path, metadata, log):
    try:
        import onnx
    except Exception:
        log("  metadata: skipped (onnx package is not installed)")
        return

    model = onnx.load(path)
    existing = {p.key: p for p in model.metadata_props}
    for key, value in metadata.items():
        prop = existing.get(key)
        if prop is None:
            prop = model.metadata_props.add()
            prop.key = key
        prop.value = "" if value is None else str(value)
    onnx.save(model, path)

    text_count = sum(1 for key in metadata.keys() if key.startswith("mcat_"))
    log(
        "  metadata: target_len={0}, len_ref={1}, len_p95={2}, text_fields={3}".format(
            metadata.get("target_len", ""),
            metadata.get("len_ref", ""),
            metadata.get("len_p95", ""),
            text_count,
        )
    )


def export_model_to_onnx(model, output_path, target_len, metadata, opset, log):
    export_model = UNet1D(in_ch=1, out_ch=1, base_ch=32)
    export_model.load_state_dict(state_dict_to_cpu(model))
    export_model.eval()

    dummy = torch.randn(1, 1, int(target_len), dtype=torch.float32)
    os.makedirs(os.path.dirname(output_path), exist_ok=True)
    torch.onnx.export(
        export_model,
        dummy,
        output_path,
        input_names=["input"],
        output_names=["output"],
        opset_version=int(opset),
        do_constant_folding=True,
    )
    add_onnx_metadata(output_path, metadata, log)


def load_first_sample_text_metadata(project_path: str):
    """
    Return metadata text from ProjectCompounds in the first ProjectSample.

    The returned dict is keyed by CompoundName so each trained .pth can store
    the text metadata for its target compound.
    """
    root = ET.parse(project_path).getroot()

    first_sample = None
    for elem in root.iter():
        if strip_ns(elem.tag) == "ProjectSample":
            first_sample = elem
            break

    if first_sample is None:
        return "", {}

    first_sample_name = get_child_text(first_sample, "Name", "") or ""
    metadata_by_compound = {}

    for comp in first_sample.iter():
        if strip_ns(comp.tag) != "ProjectCompound":
            continue

        compound_name = get_child_text(comp, "CompoundName", "")
        if not compound_name:
            continue

        metadata_by_compound[compound_name] = {
            field: get_child_text(comp, field, "") or ""
            for field in PTH_TEXT_METADATA_FIELDS
        }

    return first_sample_name, metadata_by_compound


def list_checked_compound_names(project_path: str):
    root = ET.parse(project_path).getroot()

    names = set()
    for comp in root.iter():
        if strip_ns(comp.tag) != "ProjectCompound":
            continue
        if (get_child_text(comp, "IsChecked", "false") or "false").lower() != "true":
            continue
        name = get_child_text(comp, "CompoundName", "")
        if name:
            names.add(name)
    return sorted(names)


def load_chromatograms_for_compound(
    project_path: str,
    compound_name: str,
    target_len: int = 2048,
    min_height: float = 0.0,
):
    """
    Load checked chromatograms for the given compound.
    Excludes chromatograms if max(intensity) < min_height.

    Returns:
      signals (n, L), masks (n, L), x_axis_ref (L,),
      used_n, total_n,
      rt_centers_used (list[float])  # RT_center = (RT_From+RT_To)/2 for USED chromatograms where RT valid
    """
    root = ET.parse(project_path).getroot()

    signals = []
    masks = []
    x_axis_ref = None

    total_n = 0
    used_n = 0
    rt_centers_used = []

    for comp in root.iter():
        if strip_ns(comp.tag) != "ProjectCompound":
            continue
        if (get_child_text(comp, "IsChecked", "false") or "false").lower() != "true":
            continue
        if get_child_text(comp, "CompoundName", "") != compound_name:
            continue

        total_n += 1

        x_text = get_child_text(comp, "X")
        y_text = get_child_text(comp, "Y")
        if not x_text or not y_text:
            continue

        x = np.fromstring(x_text.replace("\n", " "), sep=" ", dtype=np.float32)
        y = np.fromstring(y_text.replace("\n", " "), sep=" ", dtype=np.float32)
        if len(x) < 10 or len(y) < 10 or len(x) != len(y):
            continue

        # height filter: exclude if max intensity < min_height
        if float(np.nanmax(y)) < float(min_height):
            continue

        # RT -> mask + RT center
        rt_from_val = None
        rt_to_val = None
        try:
            rt_from_val = float(get_child_text(comp, "RT_From", "nan"))
            rt_to_val = float(get_child_text(comp, "RT_To", "nan"))
        except Exception:
            rt_from_val, rt_to_val = None, None

        mask = np.zeros_like(x, dtype=np.float32)
        if rt_from_val is not None and rt_to_val is not None:
            if np.isfinite(rt_from_val) and np.isfinite(rt_to_val) and (rt_to_val > rt_from_val):
                mask[(x >= rt_from_val) & (x <= rt_to_val)] = 1.0
                rt_centers_used.append(float(0.5 * (rt_from_val + rt_to_val)))

        x_new = np.linspace(float(x.min()), float(x.max()), target_len, dtype=np.float32)
        y_new = np.interp(x_new, x, y).astype(np.float32)
        m_new = np.interp(x_new, x, mask).astype(np.float32)
        m_new = (m_new >= 0.5).astype(np.float32)

        if x_axis_ref is None:
            x_axis_ref = x_new

        signals.append(y_new)
        masks.append(m_new)
        used_n += 1

    if used_n == 0:
        raise RuntimeError(
            f"No usable checked chromatograms found for CompoundName:\n{compound_name}\n"
            f"(min_height={min_height}, total_checked={total_n})"
        )

    return (
        np.stack(signals, axis=0),
        np.stack(masks, axis=0),
        x_axis_ref,
        used_n,
        total_n,
        rt_centers_used,
    )


class ChromDataset(Dataset):
    def __init__(self, signals: np.ndarray, masks: np.ndarray):
        self.signals = signals.astype(np.float32)
        self.masks = masks.astype(np.float32)

    def __len__(self):
        return self.signals.shape[0]

    def __getitem__(self, idx):
        x = torch.from_numpy(self.signals[idx])[None, :]
        y = torch.from_numpy(self.masks[idx])[None, :]
        return x, y


class DoubleConv1D(nn.Module):
    def __init__(self, in_ch, out_ch):
        super().__init__()
        self.net = nn.Sequential(
            nn.Conv1d(in_ch, out_ch, kernel_size=3, padding=1),
            nn.BatchNorm1d(out_ch),
            nn.ReLU(inplace=True),
            nn.Conv1d(out_ch, out_ch, kernel_size=3, padding=1),
            nn.BatchNorm1d(out_ch),
            nn.ReLU(inplace=True),
        )

    def forward(self, x):
        return self.net(x)


class UNet1D(nn.Module):
    def __init__(self, in_ch=1, out_ch=1, base_ch=32):
        super().__init__()
        self.down1 = DoubleConv1D(in_ch, base_ch)
        self.pool1 = nn.MaxPool1d(2)
        self.down2 = DoubleConv1D(base_ch, base_ch * 2)
        self.pool2 = nn.MaxPool1d(2)
        self.down3 = DoubleConv1D(base_ch * 2, base_ch * 4)
        self.pool3 = nn.MaxPool1d(2)

        self.bottleneck = DoubleConv1D(base_ch * 4, base_ch * 8)

        self.up3 = nn.ConvTranspose1d(base_ch * 8, base_ch * 4, kernel_size=2, stride=2)
        self.dec3 = DoubleConv1D(base_ch * 8, base_ch * 4)

        self.up2 = nn.ConvTranspose1d(base_ch * 4, base_ch * 2, kernel_size=2, stride=2)
        self.dec2 = DoubleConv1D(base_ch * 4, base_ch * 2)

        self.up1 = nn.ConvTranspose1d(base_ch * 2, base_ch, kernel_size=2, stride=2)
        self.dec1 = DoubleConv1D(base_ch * 2, base_ch)

        self.out_conv = nn.Conv1d(base_ch, out_ch, kernel_size=1)

    def forward(self, x):
        x1 = self.down1(x)
        x2 = self.down2(self.pool1(x1))
        x3 = self.down3(self.pool2(x2))

        xb = self.bottleneck(self.pool3(x3))

        x = self.up3(xb)
        x = torch.cat([x3, x], dim=1)
        x = self.dec3(x)

        x = self.up2(x)
        x = torch.cat([x2, x], dim=1)
        x = self.dec2(x)

        x = self.up1(x)
        x = torch.cat([x1, x], dim=1)
        x = self.dec1(x)

        return self.out_conv(x)


class ChromSegApp(tk.Tk):
    def __init__(self):
        super().__init__()
        self.title("MCAT 1D U-Net Trainer v1.3 (DirectML / ONNX)")
        self.geometry("1600x950")

        self.mcat_path = None
        self.compound_names = []

        env_device = (os.environ.get("MCAT_UNET_DEVICE", "auto") or "auto").strip().lower()
        if env_device not in ("auto", "cpu", "cuda", "directml"):
            env_device = "auto"
        self.device_mode = tk.StringVar(value=env_device)
        self.device, self.device_backend, self.initial_device_message = self._select_device(env_device, test=False)

        self.training = False
        self.out_dir = None

        # When True, next Start will reuse existing out_dir (overwrite) instead of creating a new timestamped folder.
        # This is set when Start is pressed while paused (Restart).
        self.reuse_outdir_next_run = False

        self.pause_event = threading.Event()  # set => running, clear => paused
        self.pause_event.set()

        # Stop / restart control
        self.stop_event = threading.Event()
        self.train_thread = None

        # Results gallery: 4 columns (rolling)
        self.gallery_cols = 4
        self.gallery_col_idx = 0  # 0..3

        outer = ttk.Panedwindow(self, orient="horizontal")
        outer.pack(fill=tk.BOTH, expand=True)

        self.left_frame = ttk.Frame(outer, padding=10)
        self.right_frame = ttk.Frame(outer, padding=8)
        outer.add(self.left_frame, weight=3)
        outer.add(self.right_frame, weight=7)

        right_paned = ttk.Panedwindow(self.right_frame, orient="vertical")
        right_paned.pack(fill=tk.BOTH, expand=True)

        self.plot_frame = ttk.Frame(right_paned, padding=0)
        self.log_frame = ttk.Frame(right_paned, padding=0)
        right_paned.add(self.plot_frame, weight=7)
        right_paned.add(self.log_frame, weight=3)

        self._build_left_controls(self.left_frame)
        self._build_plot_area_with_topbar(self.plot_frame)
        self._build_log(self.log_frame)

        self.update_status("App started.", write_log=True)
        self.update_status(self.initial_device_message, write_log=True)
        self.update_status(self.get_torch_environment_summary(), write_log=True)

    def _build_left_controls(self, parent):
        tk.Label(parent, text="MCAT project (.mcatproj)").pack(anchor="w")

        path_frame = tk.Frame(parent)
        path_frame.pack(fill=tk.X, pady=6)

        self.path_var = tk.StringVar(value="(no file selected)")
        tk.Entry(path_frame, textvariable=self.path_var, width=44, state="readonly").pack(side=tk.LEFT)
        tk.Button(path_frame, text="Select", command=self.select_mcatproj).pack(side=tk.LEFT, padx=6)

        tk.Label(parent, text="Compounds to train (multi-select):").pack(anchor="w", pady=(10, 0))

        comp_frame = tk.Frame(parent)
        comp_frame.pack(fill=tk.X, pady=6)

        self.comp_list = tk.Listbox(comp_frame, selectmode=tk.MULTIPLE, height=14, width=52, exportselection=False)
        comp_scroll = tk.Scrollbar(comp_frame, orient=tk.VERTICAL, command=self.comp_list.yview)
        self.comp_list.config(yscrollcommand=comp_scroll.set)
        self.comp_list.pack(side=tk.LEFT, fill=tk.BOTH, expand=True)
        comp_scroll.pack(side=tk.RIGHT, fill=tk.Y)

        btn_frame = tk.Frame(parent)
        btn_frame.pack(fill=tk.X, pady=(0, 6))
        tk.Button(btn_frame, text="All Check", command=self.select_all_compounds, width=12).pack(side=tk.LEFT)
        tk.Button(btn_frame, text="Clear", command=self.clear_all_compounds, width=12).pack(side=tk.LEFT, padx=(8, 0))

        self.comp_count_var = tk.StringVar(value="Selected: 0 / 0")
        tk.Label(parent, textvariable=self.comp_count_var, fg="gray").pack(anchor="w")
        self.comp_list.bind("<<ListboxSelect>>", lambda e: self.update_selected_count())

        tk.Label(parent, text="Device:").pack(anchor="w", pady=(12, 0))
        dev_frame = tk.Frame(parent)
        dev_frame.pack(anchor="w", pady=2)
        tk.Radiobutton(dev_frame, text="Auto", value="auto", variable=self.device_mode).pack(side=tk.LEFT)
        tk.Radiobutton(dev_frame, text="CPU", value="cpu", variable=self.device_mode).pack(side=tk.LEFT)
        tk.Radiobutton(dev_frame, text="CUDA", value="cuda", variable=self.device_mode).pack(side=tk.LEFT)
        tk.Radiobutton(dev_frame, text="DirectML", value="directml", variable=self.device_mode).pack(side=tk.LEFT)

        self.device_label_var = tk.StringVar(value=f"Current: {self.format_device_label(self.device, self.device_backend)}")
        tk.Label(parent, textvariable=self.device_label_var).pack(anchor="w")

        # manual hyper-parameters
        self.epoch_var = tk.StringVar(value="200")
        self.batch_var = tk.StringVar(value="4")
        self.lr_var = tk.StringVar(value="1e-3")
        self.use_steplr_var = tk.BooleanVar(value=True)
        self.step_size_var = tk.StringVar(value="100")
        self.gamma_var = tk.StringVar(value="0.3")
        self.len_var = tk.StringVar(value="2048")
        self.opset_var = tk.StringVar(value="17")
        self.save_pth_var = tk.BooleanVar(value=(os.environ.get("MCAT_UNET_SAVE_PTH", "0") == "1"))

        # height filter
        self.min_height_var = tk.StringVar(value="0")

        tk.Label(parent, text="Epochs:").pack(anchor="w", pady=(10, 0))
        tk.Entry(parent, textvariable=self.epoch_var, width=10).pack(anchor="w")

        tk.Label(parent, text="Batch size:").pack(anchor="w", pady=(10, 0))
        tk.Entry(parent, textvariable=self.batch_var, width=10).pack(anchor="w")

        tk.Label(parent, text="Learning rate (lr):").pack(anchor="w", pady=(10, 0))
        tk.Entry(parent, textvariable=self.lr_var, width=10).pack(anchor="w")

        tk.Label(parent, text="LR Scheduler (StepLR):").pack(anchor="w", pady=(10, 0))
        tk.Checkbutton(parent, text="Use StepLR (default ON)", variable=self.use_steplr_var).pack(anchor="w")

        steplr_frame = tk.Frame(parent)
        steplr_frame.pack(anchor="w", pady=(0, 5))
        tk.Label(steplr_frame, text="step_size").pack(side=tk.LEFT)
        tk.Entry(steplr_frame, textvariable=self.step_size_var, width=6).pack(side=tk.LEFT, padx=(4, 10))
        tk.Label(steplr_frame, text="gamma").pack(side=tk.LEFT)
        tk.Entry(steplr_frame, textvariable=self.gamma_var, width=6).pack(side=tk.LEFT, padx=(4, 0))

        tk.Label(parent, text="Resampled length (points):").pack(anchor="w", pady=(10, 0))
        tk.Entry(parent, textvariable=self.len_var, width=10).pack(anchor="w")

        tk.Label(parent, text="ONNX opset:").pack(anchor="w", pady=(10, 0))
        tk.Entry(parent, textvariable=self.opset_var, width=10).pack(anchor="w")
        tk.Checkbutton(parent, text="Also save .pth backup", variable=self.save_pth_var).pack(anchor="w")

        tk.Label(parent, text="Min chromatogram height (exclude if max < this):").pack(anchor="w", pady=(10, 0))
        tk.Entry(parent, textvariable=self.min_height_var, width=10).pack(anchor="w")


    def _build_plot_area_with_topbar(self, parent):
        topbar = tk.Frame(parent)
        topbar.pack(fill=tk.X, pady=(0, 6))

        self.btn_start = tk.Button(topbar, text="Start", fg="white", bg="green",
                                   command=self.start_training_thread, width=10)
        self.btn_start.pack(side=tk.LEFT)

        self.btn_pause = tk.Button(topbar, text="Pause", fg="black", bg="#FBBF24",
                                   command=self.toggle_pause, width=10)
        self.btn_pause.pack(side=tk.LEFT, padx=(8, 8))

        self.btn_quit = tk.Button(topbar, text="Quit", fg="white", bg="red",
                                  command=self.on_quit, width=10)
        self.btn_quit.pack(side=tk.LEFT)

        self.status_var = tk.StringVar(value="Idle")
        tk.Label(topbar, text="Status:", padx=10).pack(side=tk.LEFT)
        tk.Label(topbar, textvariable=self.status_var, fg="blue", anchor="w").pack(side=tk.LEFT, fill=tk.X, expand=True)

        # 4 rows x 4 cols gallery
        self.fig = Figure(dpi=60, constrained_layout=True)
        self.axes = [[None for _ in range(self.gallery_cols)] for _ in range(4)]
        for r in range(4):
            for c in range(self.gallery_cols):
                ax = self.fig.add_subplot(4, self.gallery_cols, r * self.gallery_cols + c + 1)
                ax.grid(True)
                self.axes[r][c] = ax

        for c in range(self.gallery_cols):
            self.axes[0][c].set_title(f"Slot {c+1}: RT", fontsize=9)
            self.axes[1][c].set_title("Chrom", fontsize=9)
            self.axes[2][c].set_title("Mask", fontsize=9)
            self.axes[3][c].set_title("Loss", fontsize=9)

        self.canvas = FigureCanvasTkAgg(self.fig, master=parent)
        self.canvas.get_tk_widget().pack(fill=tk.BOTH, expand=True)

    def _build_log(self, parent):
        tk.Label(parent, text="Log:").pack(anchor="w")
        cont = tk.Frame(parent)
        cont.pack(fill=tk.BOTH, expand=True)

        scroll = tk.Scrollbar(cont, orient=tk.VERTICAL)
        scroll.pack(side=tk.RIGHT, fill=tk.Y)

        self.log_text = tk.Text(cont, wrap="word", yscrollcommand=scroll.set)
        self.log_text.pack(side=tk.LEFT, fill=tk.BOTH, expand=True)
        scroll.config(command=self.log_text.yview)

        try:
            self.log_text.configure(font=("Consolas", 9))
        except Exception:
            pass

    def log(self, msg: str):
        ts = datetime.datetime.now().strftime("%H:%M:%S")
        line = f"[{ts}] {msg}\n"

        def _append():
            self.log_text.insert("end", line)
            self.log_text.see("end")

        self.after(0, _append)

    def update_status(self, text: str, *, write_log: bool = True):
        self.after(0, lambda: self.status_var.set(text))
        if write_log:
            self.log(text)

    def toggle_pause(self):
        if self.stop_event.is_set():
            return

        if self.pause_event.is_set():
            self.pause_event.clear()
            self.btn_pause.config(text="Resume")
            self.update_status("Paused.", write_log=True)
        else:
            self.pause_event.set()
            self.btn_pause.config(text="Pause")
            self.update_status("Resumed.", write_log=True)

    def _wait_if_paused(self):
        while self.training and (not self.pause_event.is_set()) and (not self.stop_event.is_set()):
            time.sleep(0.1)

    def get_torch_environment_summary(self):
        directml_state = "installed" if torch_directml is not None else "not installed"
        cuda_name = ""
        if torch.cuda.is_available():
            try:
                cuda_name = ", cuda_device=" + torch.cuda.get_device_name(0)
            except Exception:
                cuda_name = ""
        return (
            "Torch environment: torch={0}, cuda_build={1}, cuda_available={2}{3}, directml={4}".format(
                getattr(torch, "__version__", "?"),
                torch.version.cuda,
                torch.cuda.is_available(),
                cuda_name,
                directml_state,
            )
        )

    def format_device_label(self, device, backend):
        if backend == "cuda":
            try:
                return "CUDA: " + torch.cuda.get_device_name(0)
            except Exception:
                return "CUDA"
        if backend == "directml":
            return "DirectML: " + str(device)
        return "CPU"

    def _is_cpu_device(self):
        return self.device_backend == "cpu" or str(self.device).lower().startswith("cpu")

    def _make_directml_device(self):
        if torch_directml is None:
            detail = "torch-directml is not installed"
            if TORCH_DIRECTML_IMPORT_ERROR is not None:
                detail += ": " + str(TORCH_DIRECTML_IMPORT_ERROR)
            return None, detail
        try:
            return torch_directml.device(), None
        except Exception as exc:
            return None, str(exc)

    def _test_device(self, device, backend):
        if backend == "cpu":
            return True, ""
        try:
            test_model = nn.Sequential(
                nn.Conv1d(1, 2, kernel_size=3, padding=1),
                nn.ReLU(),
                nn.ConvTranspose1d(2, 1, kernel_size=2, stride=2),
            ).to(device)
            x = torch.randn(1, 1, 16, dtype=torch.float32).to(device)
            y = test_model(x).mean()
            y.backward()
            if backend == "cuda":
                torch.cuda.synchronize()
            return True, ""
        except Exception as exc:
            return False, str(exc)

    def _select_device(self, mode, test=True):
        mode = (mode or "auto").lower()

        def choose_cpu(message):
            return torch.device("cpu"), "cpu", message

        def choose_cuda(prefix):
            if not torch.cuda.is_available():
                return None, None, "CUDA is not available"
            device = torch.device("cuda")
            if test:
                ok, detail = self._test_device(device, "cuda")
                if not ok:
                    return None, None, "CUDA test failed: " + detail
            return device, "cuda", prefix + "Using CUDA."

        def choose_directml(prefix):
            device, detail = self._make_directml_device()
            if device is None:
                return None, None, "DirectML is not available: " + detail
            if test:
                ok, test_detail = self._test_device(device, "directml")
                if not ok:
                    return None, None, "DirectML test failed: " + test_detail
            return device, "directml", prefix + "Using DirectML."

        if mode == "cpu":
            return choose_cpu("Using CPU.")

        if mode == "cuda":
            device, backend, message = choose_cuda("")
            if device is None:
                return choose_cpu(message + " Falling back to CPU.")
            return device, backend, message

        if mode == "directml":
            device, backend, message = choose_directml("")
            if device is None:
                return choose_cpu(message + " Falling back to CPU.")
            return device, backend, message

        device, backend, message = choose_cuda("Auto: ")
        if device is not None:
            return device, backend, message

        cuda_message = message
        device, backend, message = choose_directml("Auto: ")
        if device is not None:
            return device, backend, message

        return choose_cpu("Auto: " + cuda_message + " DirectML is not available. Using CPU.")

    def switch_to_cpu_after_device_error(self, detail):
        failed = self.format_device_label(self.device, self.device_backend)
        self.device = torch.device("cpu")
        self.device_backend = "cpu"
        self.after(0, lambda: self.device_label_var.set("Current: CPU"))
        self.update_status(
            "Device error on {0}; retrying this compound on CPU. Detail: {1}".format(failed, detail),
            write_log=True,
        )

    def resolve_device(self):
        mode = self.device_mode.get()
        self.device, self.device_backend, msg = self._select_device(mode, test=True)
        if self.device_backend == "cpu" and mode in ("cuda", "directml"):
            messagebox.showwarning("Warning", msg)

        self.device_label_var.set(f"Current: {self.format_device_label(self.device, self.device_backend)}")
        self.update_status(msg, write_log=True)

    def set_compound_list(self, names):
        self.comp_list.delete(0, "end")
        for n in names:
            self.comp_list.insert("end", n)
        self.select_all_compounds()

    def select_all_compounds(self):
        if self.comp_list.size() > 0:
            self.comp_list.select_set(0, "end")
        self.update_selected_count()

    def clear_all_compounds(self):
        self.comp_list.select_clear(0, "end")
        self.update_selected_count()

    def get_selected_compounds(self):
        return [self.comp_list.get(i) for i in self.comp_list.curselection()]

    def update_selected_count(self):
        total = len(self.compound_names) if self.compound_names else 0
        sel = len(self.get_selected_compounds()) if total else 0
        self.comp_count_var.set(f"Selected: {sel} / {total}")

    def select_mcatproj(self):
        path = filedialog.askopenfilename(
            title="Select MCAT project file",
            filetypes=[("MCAT Project", "*.mcatproj"), ("All files", "*.*")]
        )
        if not path:
            return

        self.mcat_path = path
        self.path_var.set(path)
        self.update_status("Reading compound list...", write_log=True)

        try:
            names = list_checked_compound_names(path)
        except Exception as e:
            messagebox.showerror("Error", f"Failed to read project:\n{e}")
            self.update_status("Error while reading project.", write_log=True)
            return

        self.compound_names = names
        self.set_compound_list(names)
        if not names:
            messagebox.showwarning("Warning", "No ProjectCompound with IsChecked=true found.")
            self.update_status("No checked compounds.", write_log=True)
        else:
            self.update_status(f"Loaded {len(names)} checked compounds. (All selected)", write_log=True)

    def on_quit(self):
        if self.training and (not messagebox.askokcancel("Confirm", "Training is running. Really quit?")):
            return
        # request stop so background thread can exit quickly
        self.stop_event.set()
        self.pause_event.set()
        self.training = False
        self.destroy()

    def start_training_thread(self):
        # If training is running and currently paused, Start => restart from scratch
        if self.training:
            if not self.pause_event.is_set():
                # paused -> restart
                self.update_status("Restart requested: stopping current training...", write_log=True)
                # Reuse the same out_dir on restart (overwrite)
                self.reuse_outdir_next_run = True
                self.stop_event.set()
                self.pause_event.set()
                self.btn_pause.config(text="Pause")
                threading.Thread(target=self._restart_after_stop, daemon=True).start()
            else:
                messagebox.showinfo("Info", "Training is already running.")
            return

        # normal start -> new out_dir
        self.reuse_outdir_next_run = False
        self.stop_event.clear()
        self.pause_event.set()
        self.btn_pause.config(text="Pause")

        t = threading.Thread(target=self._run_training_safe, daemon=True)
        self.train_thread = t
        t.start()

    def _restart_after_stop(self):
        """Wait for the current training thread to stop, then start a fresh training run.

        This is triggered when the user presses Start during Pause.
        """
        t = self.train_thread
        if t is not None and t.is_alive():
            # wait (non-blocking UI)
            for _ in range(300):  # up to ~30s in 0.1s steps; typical stop is much faster
                if not t.is_alive():
                    break
                time.sleep(0.1)

        # Start fresh
        self.stop_event.clear()
        self.pause_event.set()
        self.btn_pause.config(text="Pause")
        self.update_status("Restarting training from scratch...", write_log=True)

        nt = threading.Thread(target=self._run_training_safe, daemon=True)
        self.train_thread = nt
        nt.start()

    def _run_training_safe(self):
        try:
            # ensure stop flag is clear for a fresh run
            self.stop_event.clear()
            reuse = bool(self.reuse_outdir_next_run)
            # reset immediately so a later normal Start doesn't unexpectedly reuse
            self.reuse_outdir_next_run = False
            self.run_training(reuse_outdir=reuse)
        except Exception as e:
            self.training = False
            self.update_status(f"Training crashed: {e}", write_log=True)
            self.status_var.set("Idle")

    def plot_rt_boxplot_for_compound(self, rt_centers_used, compound, used_n, total_n, min_height, col):
        def _plot():
            ax_rt = self.axes[0][col]
            ax_chr = self.axes[1][col]
            ax_m = self.axes[2][col]
            ax_l = self.axes[3][col]

            ax_rt.clear()
            if rt_centers_used and len(rt_centers_used) > 0:
                ax_rt.boxplot(rt_centers_used, vert=False, showfliers=True)
                ax_rt.set_title(f"{compound}\nRT box (train-only)", fontsize=9)
                ax_rt.set_xlabel("RT center", fontsize=8)
                ax_rt.grid(True)
            else:
                ax_rt.text(0.5, 0.5, "No valid RT\n(train-only)", ha="center", va="center", transform=ax_rt.transAxes)
                ax_rt.set_title(f"{compound}\nRT box (train-only)", fontsize=9)
                ax_rt.set_axis_off()

            ax_chr.clear()
            ax_chr.text(0.5, 0.5, "Training...", ha="center", va="center", transform=ax_chr.transAxes)
            ax_chr.set_axis_off()

            ax_m.clear()
            ax_m.text(0.5, 0.5, "Training...", ha="center", va="center", transform=ax_m.transAxes)
            ax_m.set_axis_off()

            ax_l.clear()
            ax_l.text(0.5, 0.5, "Training...", ha="center", va="center", transform=ax_l.transAxes)
            ax_l.set_axis_off()

            self.canvas.draw()

        self.after(0, _plot)

        valid_rt = len(rt_centers_used) if rt_centers_used is not None else 0
        self.update_status(
            f"[RT] slot={col+1} {compound} | train chrom: {used_n}/{total_n} (min_height={min_height}) | RT_valid={valid_rt}",
            write_log=True
        )

    def train_model_for_compound(
        self,
        signals,
        masks,
        epochs,
        batch_size,
        lr,
        use_steplr,
        step_size,
        gamma,
        index,
        total,
        compound,
    ):
        dataset = ChromDataset(signals, masks)
        bs = int(max(1, min(batch_size, len(dataset))))
        loader = DataLoader(dataset, batch_size=bs, shuffle=True)

        model = UNet1D(in_ch=1, out_ch=1, base_ch=32).to(self.device)
        criterion = nn.BCEWithLogitsLoss()
        optimizer = torch.optim.Adam(model.parameters(), lr=lr)

        scheduler = None
        if use_steplr:
            scheduler = torch.optim.lr_scheduler.StepLR(optimizer, step_size=step_size, gamma=gamma)

        losses = []
        n_samples = len(dataset)

        model.train()
        for ep in range(1, epochs + 1):
            self._wait_if_paused()
            if self.stop_event.is_set():
                break

            epoch_loss = 0.0
            for x, y in loader:
                self._wait_if_paused()
                if self.stop_event.is_set():
                    break

                x = x.to(self.device)
                y = y.to(self.device)

                optimizer.zero_grad()
                logits = model(x)
                loss = criterion(logits, y)
                loss.backward()
                optimizer.step()

                epoch_loss += loss.item() * x.size(0)

            if self.stop_event.is_set():
                break

            epoch_loss /= max(1, n_samples)
            losses.append(epoch_loss)

            if scheduler is not None:
                scheduler.step()

            cur_lr = optimizer.param_groups[0].get("lr", lr)
            self.update_status(
                f"[{index}/{total}] {compound} | Epoch {ep}/{epochs} | loss={epoch_loss:.4f} | lr={cur_lr:.3g}",
                write_log=False
            )

        final_loss = float(losses[-1]) if losses else float("nan")
        self.log(f"[{index}/{total}] {compound} | FINAL loss={final_loss:.6f}")

        model.eval()
        with torch.no_grad():
            sig0 = signals[0]
            x0 = torch.from_numpy(sig0[None, None, :]).to(self.device)
            pred_logits = model(x0)
            pred_mask = torch.sigmoid(pred_logits).cpu().numpy()[0, 0, :]

        return bs, model, losses, final_loss, pred_mask

    def run_training(self, reuse_outdir: bool = False):
        if not self.mcat_path or not os.path.isfile(self.mcat_path):
            messagebox.showerror("Error", "Please select a valid .mcatproj file.")
            return

        selected = self.get_selected_compounds()
        if not selected:
            messagebox.showerror("Error", "No compounds selected. Please select at least one compound to train.")
            return

        try:
            epochs = int(self.epoch_var.get())
            batch_size = int(self.batch_var.get())
            target_len = int(self.len_var.get())
            opset = int(self.opset_var.get())
            lr = float(self.lr_var.get())

            use_steplr = bool(self.use_steplr_var.get())
            step_size = int(self.step_size_var.get())
            gamma = float(self.gamma_var.get())

            min_height = float(self.min_height_var.get())
        except ValueError:
            messagebox.showerror("Error", "Epochs/batch/length/opset/step_size must be integers; lr/gamma/min_height must be numbers.")
            return

        if epochs <= 0 or batch_size <= 0 or target_len <= 0 or lr <= 0:
            messagebox.showerror("Error", "Epochs, batch size, length, lr must be > 0.")
            return
        if opset <= 0:
            messagebox.showerror("Error", "ONNX opset must be > 0.")
            return
        if use_steplr and (step_size <= 0 or not (0 < gamma < 1.0)):
            messagebox.showerror("Error", "StepLR: step_size must be >0 and gamma must be (0,1).")
            return
        if min_height < 0:
            messagebox.showerror("Error", "min_height must be >= 0.")
            return

        self.resolve_device()

        project_dir = os.path.dirname(self.mcat_path)

        save_pth = bool(self.save_pth_var.get())

        # Output folder policy:
        # - Normal Start: create a new timestamped folder (keeps past runs)
        # - Restart (Start pressed while paused): reuse the SAME folder.
        #   NOTE: On restart we DO NOT delete existing unet1d_*.onnx/.pth. New training results
        #   will overwrite files only when the same compound name produces the same file name.
        if reuse_outdir and self.out_dir and os.path.isdir(self.out_dir):
            # Reuse existing out_dir
            pass
        else:
            ts = datetime.datetime.now().strftime("%Y%m%d_%H%M%S")
            self.out_dir = os.path.join(project_dir, f"unet1d_train_{ts}")
            os.makedirs(self.out_dir, exist_ok=True)

        self.training = True
        self.stop_event.clear()
        total = len(selected)
        trained_count = 0

        self.update_status(f"Output folder: {self.out_dir}", write_log=True)
        self.update_status(f"Training start. Selected compounds: {total}", write_log=True)

        first_sample_name = ""
        first_sample_metadata = {}
        try:
            first_sample_name, first_sample_metadata = load_first_sample_text_metadata(self.mcat_path)
            if first_sample_name:
                self.update_status(
                    f"Loaded text metadata from first sample: {first_sample_name}",
                    write_log=True
                )
        except Exception as e:
            self.update_status(f"Could not read text metadata from first sample: {e}", write_log=True)

        for i, compound in enumerate(selected, start=1):
            self._wait_if_paused()
            if self.stop_event.is_set():
                break

            col = self.gallery_col_idx
            self.update_status(f"[{i}/{total}] Start: {compound} (slot {col+1})", write_log=True)

            try:
                signals, masks, x_axis, used_n, total_n, rt_centers_used = load_chromatograms_for_compound(
                    self.mcat_path,
                    compound_name=compound,
                    target_len=target_len,
                    min_height=min_height
                )
            except Exception as e:
                self.update_status(f"[{i}/{total}] Skip '{compound}' (load error): {e}", write_log=True)
                continue

            self.update_status(
                f"[{i}/{total}] {compound} | training chromatograms: {used_n}/{total_n} (min_height={min_height})",
                write_log=True
            )

            self.plot_rt_boxplot_for_compound(rt_centers_used, compound, used_n, total_n, min_height, col)

            try:
                bs, model, losses, final_loss, pred_mask = self.train_model_for_compound(
                    signals,
                    masks,
                    epochs,
                    batch_size,
                    lr,
                    use_steplr,
                    step_size,
                    gamma,
                    i,
                    total,
                    compound,
                )
            except Exception as e:
                if self._is_cpu_device():
                    self.update_status(f"[{i}/{total}] Skip '{compound}' (training error): {e}", write_log=True)
                    continue

                self.switch_to_cpu_after_device_error(e)
                try:
                    bs, model, losses, final_loss, pred_mask = self.train_model_for_compound(
                        signals,
                        masks,
                        epochs,
                        batch_size,
                        lr,
                        use_steplr,
                        step_size,
                        gamma,
                        i,
                        total,
                        compound,
                    )
                except Exception as cpu_error:
                    self.update_status(f"[{i}/{total}] Skip '{compound}' (CPU retry failed): {cpu_error}", write_log=True)
                    continue

            if self.stop_event.is_set():
                break

            sig0 = signals[0]
            mask0 = masks[0]
            self.plot_results_to_slot(x_axis, sig0, mask0, pred_mask, losses, compound, col)

            try:
                safe_name = safe_compound_name(compound)
                out_path = os.path.join(self.out_dir, f"unet1d_{safe_name}.onnx")
                pth_path = os.path.join(self.out_dir, f"unet1d_{safe_name}.pth")
                compound_metadata, metadata_match_name = find_compound_metadata(compound, first_sample_metadata)
                if metadata_match_name:
                    self.update_status(
                        f"[{i}/{total}] {compound} | text metadata matched: {metadata_match_name}",
                        write_log=True
                    )
                else:
                    self.update_status(
                        f"[{i}/{total}] {compound} | text metadata not found.",
                        write_log=True
                    )

                metadata = {
                    "version": "unet1d_directml_onnx_v1_3",
                    "compound_name": compound,
                    "target_len": int(target_len),
                    "len_ref": 0,
                    "len_p95": 0,
                    "peak_count": 1,
                    "train_peak_count_mode": 1,
                    "train_peak_count_min": 1,
                    "train_peak_count_max": 1,
                    "peak_width_points_median": 0,
                    "train_epochs": int(epochs),
                    "train_batch_size": int(bs),
                    "train_lr_init": float(lr),
                    "train_device_backend": self.device_backend,
                    "min_height": float(min_height),
                    "use_steplr": bool(use_steplr),
                    "steplr_step_size": int(step_size) if use_steplr else None,
                    "steplr_gamma": float(gamma) if use_steplr else None,
                    "final_loss": float(final_loss),
                    "mcat_first_sample_name": first_sample_name,
                }
                if isinstance(compound_metadata, dict):
                    for key in PTH_TEXT_METADATA_FIELDS:
                        value = compound_metadata.get(key, "")
                        if value is not None:
                            metadata["mcat_" + key] = str(value)
                    metadata["mcat_first_sample_project_compound_metadata"] = json.dumps(
                        compound_metadata,
                        ensure_ascii=False,
                    )

                export_model_to_onnx(model, out_path, target_len, metadata, opset, self.log)
                saved_paths = [out_path]

                if save_pth:
                    ckpt = {
                        "version": "unet1d_directml_onnx_v1_3",
                        "compound_name": compound,
                        "model_state_dict": state_dict_to_cpu(model),
                        "first_sample_name": first_sample_name,
                        "first_sample_project_compound_metadata": compound_metadata,
                        "target_len": int(target_len),
                        "train_epochs": int(epochs),
                        "train_batch_size": int(bs),
                        "train_lr_init": float(lr),
                        "train_device_backend": self.device_backend,
                        "min_height": float(min_height),
                        "use_steplr": bool(use_steplr),
                        "steplr_step_size": int(step_size) if use_steplr else None,
                        "steplr_gamma": float(gamma) if use_steplr else None,
                        "final_loss": float(final_loss),
                    }
                    torch.save(ckpt, pth_path)
                    saved_paths.append(pth_path)

                trained_count += 1
                self.update_status(f"[{i}/{total}] Done: {compound} | saved: {', '.join(saved_paths)}", write_log=True)
            except Exception as e:
                self.update_status(f"[{i}/{total}] Done: {compound} | save failed: {e}", write_log=True)

            self.gallery_col_idx = (self.gallery_col_idx + 1) % self.gallery_cols

        self.training = False
        self.train_thread = None
        if self.stop_event.is_set():
            self.update_status(f"Stopped. Trained {trained_count}/{total} compounds. Folder: {self.out_dir}", write_log=True)
        else:
            self.update_status(f"All done. Trained {trained_count}/{total} compounds. Folder: {self.out_dir}", write_log=True)
        self.status_var.set("Idle")

    def plot_results_to_slot(self, x_axis, signal, mask_true, mask_pred, losses, compound, col):
        def _plot():
            ax_chr = self.axes[1][col]
            ax_m = self.axes[2][col]
            ax_l = self.axes[3][col]

            ax_chr.clear()
            ax_m.clear()
            ax_l.clear()

            ax_chr.plot(x_axis, signal)
            ax_chr.set_title(f"{compound}\nChromatogram", fontsize=9)
            ax_chr.set_xlabel("Time", fontsize=8)
            ax_chr.set_ylabel("I", fontsize=8)
            ax_chr.grid(True)

            ax_m.plot(x_axis, mask_true, label="True")
            ax_m.plot(x_axis, mask_pred, label="Pred", alpha=0.7)
            ax_m.set_title("Mask", fontsize=9)
            ax_m.set_xlabel("Time", fontsize=8)
            ax_m.set_ylabel("V", fontsize=8)
            ax_m.legend(fontsize=7)
            ax_m.grid(True)

            ax_l.plot(range(1, len(losses) + 1), losses)
            ax_l.set_title("Loss", fontsize=9)
            ax_l.set_xlabel("Ep", fontsize=8)
            ax_l.set_ylabel("L", fontsize=8)
            ax_l.grid(True)

            self.canvas.draw()

        self.after(0, _plot)


if __name__ == "__main__":
    app = ChromSegApp()
    app.mainloop()
