# -*- coding: utf-8 -*-
import argparse
import json
import os
import re
import sys

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


def log(message):
    print(message, flush=True)


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


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):
    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)


class ChromDataset(Dataset):
    def __init__(self, signals, masks):
        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)


def select_device(mode):
    mode = (mode or "auto").strip().lower()

    if mode in ("auto", "cuda") and torch.cuda.is_available():
        return torch.device("cuda"), "cuda"

    if mode in ("auto", "directml"):
        try:
            import torch_directml
            return torch_directml.device(), "directml"
        except Exception as exc:
            if mode == "directml":
                log("directml unavailable, using cpu: {0}".format(exc))

    return torch.device("cpu"), "cpu"


def load_input(path):
    with open(path, "r", encoding="utf-8-sig") as f:
        return json.load(f)


def resample_record(record, target_len):
    x = np.asarray(record.get("x") or [], dtype=np.float32)
    y = np.asarray(record.get("y") or [], dtype=np.float32)
    if x.size < 10 or y.size < 10 or x.size != y.size:
        raise ValueError("invalid x/y length")

    order = np.argsort(x)
    x = x[order]
    y = y[order]

    mask = np.zeros_like(x, dtype=np.float32)
    for pair in record.get("ranges") or []:
        if not isinstance(pair, (list, tuple)) or len(pair) < 2:
            continue
        try:
            start = float(pair[0])
            end = float(pair[1])
        except Exception:
            continue
        if not np.isfinite(start) or not np.isfinite(end):
            continue
        if end < start:
            start, end = end, start
        if end > start:
            mask[(x >= start) & (x <= end)] = 1.0

    if float(mask.max()) <= 0.0:
        raise ValueError("no teacher range")

    x_new = np.linspace(float(x.min()), float(x.max()), int(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)
    return y_new, m_new


def count_valid_ranges(record):
    count = 0
    for pair in record.get("ranges") or []:
        if not isinstance(pair, (list, tuple)) or len(pair) < 2:
            continue
        try:
            start = float(pair[0])
            end = float(pair[1])
        except Exception:
            continue
        if not np.isfinite(start) or not np.isfinite(end):
            continue
        if end != start:
            count += 1
    return count


def choose_peak_count(range_counts):
    counts = [int(c) for c in range_counts if int(c) > 0]
    if not counts:
        return 1

    freq = {}
    for count in counts:
        freq[count] = freq.get(count, 0) + 1

    return sorted(freq.items(), key=lambda item: (-item[1], item[0]))[0][0]


def mask_segment_widths(mask):
    widths = []
    i = 0
    while i < int(mask.size):
        while i < int(mask.size) and float(mask[i]) < 0.5:
            i += 1
        if i >= int(mask.size):
            break
        start = i
        while i < int(mask.size) and float(mask[i]) >= 0.5:
            i += 1
        widths.append(i - start)
    return widths


def export_model_to_onnx(model, output_path, target_len, metadata, opset):
    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)


def train_one_prepared(compound, records, args, device, device_backend, signals, masks, range_counts, segment_widths, skipped):
    signals = np.stack(signals, axis=0)
    masks = np.stack(masks, axis=0)
    dataset = ChromDataset(signals, masks)
    batch_size = int(max(1, min(args.batch_size, len(dataset))))
    loader = DataLoader(dataset, batch_size=batch_size, shuffle=True)

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

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

    log(
        "train start: {0}, records={1}, skipped={2}, epochs={3}, device={4}".format(
            compound,
            len(dataset),
            skipped,
            args.epochs,
            device_backend,
        )
    )

    final_loss = float("nan")
    for epoch in range(1, int(args.epochs) + 1):
        model.train()
        total_loss = 0.0
        total_count = 0

        for x, y in loader:
            x = x.to(device)
            y = y.to(device)

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

            total_loss += float(loss.item()) * int(x.size(0))
            total_count += int(x.size(0))

        if scheduler is not None:
            scheduler.step()

        final_loss = total_loss / max(1, total_count)
        if epoch == 1 or epoch == int(args.epochs) or epoch % max(1, int(args.epochs) // 5) == 0:
            log("train epoch: {0}, {1}/{2}, loss={3:.6f}".format(compound, epoch, args.epochs, final_loss))

    first_metadata = records[0].get("metadata") or {}
    peak_count = choose_peak_count(range_counts)
    peak_width_points = int(np.median(np.asarray(segment_widths, dtype=np.float32))) if segment_widths else 0
    metadata = {
        "version": "unet1d_wpf_headless_v1",
        "compound_name": compound,
        "target_len": int(args.target_len),
        "len_ref": 0,
        "len_p95": 0,
        "peak_count": int(peak_count),
        "train_peak_count_mode": int(peak_count),
        "train_peak_count_min": int(min(range_counts)) if range_counts else 0,
        "train_peak_count_max": int(max(range_counts)) if range_counts else 0,
        "peak_width_points_median": int(peak_width_points),
        "train_epochs": int(args.epochs),
        "train_batch_size": int(batch_size),
        "train_lr_init": float(args.lr),
        "train_device_backend": device_backend,
        "final_loss": float(final_loss),
        "train_records": int(len(dataset)),
    }
    for key, value in first_metadata.items():
        if re.match(r"^[A-Za-z0-9_]+$", str(key)):
            metadata["mcat_" + str(key)] = "" if value is None else str(value)

    output_path = os.path.join(args.output_dir, "unet1d_{0}.onnx".format(safe_compound_name(compound)))
    export_model_to_onnx(model, output_path, args.target_len, metadata, args.opset)
    log("saved: {0}".format(output_path))
    return True


def train_one(compound, records, args, device, device_backend):
    signals = []
    masks = []
    range_counts = []
    segment_widths = []
    skipped = 0

    for record in records:
        try:
            sig, mask = resample_record(record, args.target_len)
        except Exception as exc:
            skipped += 1
            log("skip record: {0}: {1}".format(compound, exc))
            continue
        signals.append(sig)
        masks.append(mask)
        range_counts.append(count_valid_ranges(record))
        segment_widths.extend(mask_segment_widths(mask))

    if not signals:
        log("skip compound: {0}: no usable records".format(compound))
        return False

    try:
        return train_one_prepared(compound, records, args, device, device_backend, signals, masks, range_counts, segment_widths, skipped)
    except Exception as exc:
        if device_backend == "cpu":
            raise

        log("training failed on {0}: {1}".format(device_backend, exc))
        log("retrying on cpu: {0}".format(compound))
        return train_one_prepared(
            compound,
            records,
            args,
            torch.device("cpu"),
            "cpu-retry",
            signals,
            masks,
            range_counts,
            segment_widths,
            skipped,
        )


def main(argv=None):
    parser = argparse.ArgumentParser()
    parser.add_argument("--input", required=True)
    parser.add_argument("--output-dir", required=True)
    parser.add_argument("--epochs", type=int, default=200)
    parser.add_argument("--batch-size", type=int, default=4)
    parser.add_argument("--target-len", type=int, default=2048)
    parser.add_argument("--lr", type=float, default=1.0e-3)
    parser.add_argument("--opset", type=int, default=17)
    parser.add_argument("--device", default=os.environ.get("MCAT_UNET_DEVICE", "auto"))
    parser.add_argument("--use-steplr", action="store_true")
    parser.add_argument("--step-size", type=int, default=100)
    parser.add_argument("--gamma", type=float, default=0.3)
    args = parser.parse_args(argv)

    payload = load_input(args.input)
    records = payload.get("records") or []
    if not records:
        raise RuntimeError("input has no records")

    os.makedirs(args.output_dir, exist_ok=True)
    device, backend = select_device(args.device)
    log("device: {0}".format(backend))

    by_compound = {}
    for record in records:
        compound = (record.get("compound_name") or "").strip()
        if not compound:
            continue
        by_compound.setdefault(compound, []).append(record)

    saved = 0
    for compound in sorted(by_compound.keys(), key=lambda s: s.lower()):
        if train_one(compound, by_compound[compound], args, device, backend):
            saved += 1

    log("finished: saved={0}, compounds={1}".format(saved, len(by_compound)))
    if saved == 0:
        return 2
    return 0


if __name__ == "__main__":
    try:
        sys.exit(main())
    except Exception as exc:
        log("error: {0}".format(exc))
        sys.exit(1)
