#!/usr/bin/env python3
# -*- coding: utf-8 -*-

from __future__ import annotations

import csv
from pathlib import Path
from typing import Optional, Tuple, List

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


STEP2_DIR = Path("./step2"); STEP2_DIR.mkdir(parents=True, exist_ok=True)


class TableDataset(Dataset):
    def __init__(self, x: torch.Tensor, y: Optional[torch.Tensor] = None):
        self.x = x.to(torch.float32)
        self.y = None if y is None else y.to(torch.long)

    def __len__(self) -> int:
        return int(self.x.shape[0])

    def __getitem__(self, idx: int):
        if self.y is None:
            return self.x[idx]
        return self.x[idx], self.y[idx]


def _load_from_csv() -> Optional[Tuple[torch.Tensor, torch.Tensor, torch.Tensor]]:
    td, tl, te = STEP2_DIR / "train_data.csv", STEP2_DIR / "train_label.csv", STEP2_DIR / "test_data.csv"
    if not (td.exists() and tl.exists() and te.exists()):
        return None

    def _read_features_csv(path: Path) -> List[List[float]]:
        rows: List[List[float]] = []
        with path.open("r", newline="") as f:
            reader = csv.reader(f)
            header = True
            for row in reader:
                if header:
                    try:
                        [float(v) for v in row]
                    except Exception:
                        header = False
                        continue
                    header = False
                try:
                    rows.append([float(v) for v in row])
                except Exception:
                    continue
        return rows

    xtr = torch.tensor(_read_features_csv(td), dtype=torch.float32)
    xte = torch.tensor(_read_features_csv(te), dtype=torch.float32)

    targets: List[int] = []
    with tl.open("r", newline="") as f:
        rdr = csv.DictReader(f)
        if not rdr.fieldnames or "target" not in rdr.fieldnames:
            f.seek(0)
            for row in csv.reader(f):
                try:
                    targets.append(int(row[0]))
                except Exception:
                    continue
        else:
            for r in rdr:
                targets.append(int(r["target"]))
    ytr = torch.tensor(targets, dtype=torch.long)
    assert xtr.shape[0] == ytr.shape[0]
    return xtr, ytr, xte


def _load_from_sklearn() -> Optional[Tuple[torch.Tensor, torch.Tensor, torch.Tensor]]:
    try:
        from sklearn.datasets import load_iris
        from sklearn.model_selection import train_test_split
    except Exception:
        return None
    iris = load_iris(); X, y = iris["data"], iris["target"]
    xtr, xte, ytr, _ = train_test_split(X, y, test_size=0.2, stratify=y, random_state=42)
    return torch.tensor(xtr, dtype=torch.float32), torch.tensor(ytr, dtype=torch.long), torch.tensor(xte, dtype=torch.float32)


class IrisMLP(nn.Module):
    def __init__(self, in_dim: int = 4, hidden: int = 16, num_classes: int = 3):
        super().__init__()
        self.net = nn.Sequential(
            nn.Linear(in_dim, hidden), nn.ReLU(),
            nn.Linear(hidden, hidden), nn.ReLU(),
            nn.Linear(hidden, num_classes)
        )

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        return self.net(x)


def train_and_predict(epochs: int = 50, batch_size: int = 32, lr: float = 1e-2, weight_decay: float = 0.0, device: Optional[str] = None) -> None:
    device = device or ("cuda" if torch.cuda.is_available() else "cpu")
    data = _load_from_csv() or _load_from_sklearn()
    if data is None:
        print("[Iris-MLP] 数据缺失：准备 ./step2/*.csv 或安装 scikit-learn。")
        return

    xtr, ytr, xte = data
    train_ds = TableDataset(xtr, ytr)
    test_ds = TableDataset(xte, None)
    train_loader = DataLoader(train_ds, batch_size=batch_size, shuffle=True)

    model = IrisMLP(in_dim=int(train_ds.x.shape[1]), num_classes=3).to(device)
    crit = nn.CrossEntropyLoss()
    opt = torch.optim.Adam(model.parameters(), lr=lr, weight_decay=weight_decay)

    model.train()
    for ep in range(1, epochs + 1):
        loss_sum = 0.0; correct = 0; total = 0
        for xb, yb in train_loader:
            xb, yb = xb.to(device), yb.to(device)
            logits = model(xb)
            loss = crit(logits, yb)
            opt.zero_grad(); loss.backward(); opt.step()
            loss_sum += float(loss.item()) * xb.size(0)
            correct += int((logits.argmax(1) == yb).sum().item())
            total += int(xb.size(0))
        print(f"[Iris-MLP] Epoch {ep:03d} | Loss {loss_sum/total:.4f} | Acc {correct/total:.4f}")

    model.eval()
    with torch.no_grad():
        preds = model(test_ds.x.to(device)).argmax(1).cpu().tolist()

    torch.save(model.state_dict(), STEP2_DIR / "iris_mlp.pt")
    out_csv = STEP2_DIR / "result.csv"
    with out_csv.open("w", newline="") as f:
        w = csv.writer(f); w.writerow(["target"]); [w.writerow([int(p)]) for p in preds]
    print(f"[Iris-MLP] saved: {out_csv}")


if __name__ == "__main__":
    train_and_predict()