"""Train a task-specific DISTAL predictor and predict candidate formulas."""

from __future__ import annotations

import argparse
import copy
import random
from pathlib import Path

import numpy as np
import pandas as pd
import torch
from matminer.featurizers import composition as cf
from matminer.featurizers.base import MultipleFeaturizer
from pymatgen.core import Composition
from sklearn.metrics import mean_absolute_error
from sklearn.model_selection import KFold
from sklearn.preprocessing import StandardScaler
from torch import nn
from torch.optim.lr_scheduler import ReduceLROnPlateau
from torch.utils.data import DataLoader, TensorDataset


SCRIPT_DIR = Path(__file__).resolve().parent
DEFAULT_MODEL_DIR = SCRIPT_DIR / "models"
DEFAULT_FUSION_WEIGHTS = DEFAULT_MODEL_DIR / "adaptive_denoising_resnet.pth"
DEFAULT_STUDENT_WEIGHTS = DEFAULT_MODEL_DIR / "student_encoder_best.pth"
DEFAULT_OUTPUT_CSV = SCRIPT_DIR / "outputs" / "distal_predictions.csv"
DEFAULT_BUNDLE_PATH = SCRIPT_DIR / "outputs" / "distal_task_bundle.pt"
DEVICE = "cuda" if torch.cuda.is_available() else "cpu"


class Config:
    input_dim = 145
    latent_dim = 72
    structure_dim = 256
    min_delta = 1e-4
    num_workers = 0


class ResBlock(nn.Module):
    def __init__(self, dim: int):
        super().__init__()
        self.block = nn.Sequential(
            nn.Linear(dim, dim),
            nn.BatchNorm1d(dim),
            nn.LeakyReLU(0.2),
        )

    def forward(self, x):
        return x + self.block(x)


class ADAE_ResNet(nn.Module):
    def __init__(self, input_dim: int = 145, latent_dim: int = 72):
        super().__init__()
        self.encoder = nn.Sequential(
            nn.Linear(input_dim, 128),
            ResBlock(128),
            nn.Linear(128, latent_dim),
        )
        self.decoder = nn.Sequential(
            nn.Linear(latent_dim, 128),
            ResBlock(128),
            nn.Linear(128, input_dim),
        )
        self.noise_factor = nn.Parameter(torch.tensor(0.1))

    def forward(self, x, epoch_ratio: float = 0.5):
        noisy_x = x
        if self.training:
            noise = 0.01 + 0.49 * torch.sigmoid(self.noise_factor)
            noisy_x = x + torch.randn_like(x) * noise
        z = self.encoder(noisy_x)
        xp = self.decoder(z)
        return xp, z, {}


class StudentEncoder(nn.Module):
    def __init__(self, input_dim: int = 145, output_dim: int = 256):
        super().__init__()
        self.net = nn.Sequential(
            nn.Linear(input_dim, 512),
            nn.BatchNorm1d(512),
            nn.SiLU(),
            nn.Dropout(0.2),
            nn.Linear(512, 512),
            nn.BatchNorm1d(512),
            nn.SiLU(),
            nn.Dropout(0.2),
            nn.Linear(512, output_dim),
        )

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


class PropertyPredictor(nn.Module):
    def __init__(self, input_dim: int):
        super().__init__()
        hidden = 512
        self.backbone = nn.Sequential(
            nn.Linear(input_dim, hidden),
            nn.BatchNorm1d(hidden),
            nn.GELU(),
            nn.Dropout(0.15),
            nn.Linear(hidden, hidden // 2),
            nn.BatchNorm1d(hidden // 2),
            nn.GELU(),
            nn.Dropout(0.10),
            nn.Linear(hidden // 2, hidden // 4),
            nn.GELU(),
            nn.Dropout(0.05),
            nn.Linear(hidden // 4, 1),
        )

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


class TargetScaler:
    def __init__(self):
        self.mean_ = 0.0
        self.std_ = 1.0

    def fit(self, y):
        y = np.asarray(y, dtype=np.float32)
        self.mean_ = float(np.mean(y))
        self.std_ = float(np.std(y))
        if self.std_ < 1e-8:
            self.std_ = 1.0
        return self

    def transform(self, y):
        return (np.asarray(y, dtype=np.float32) - self.mean_) / self.std_

    def inverse_transform(self, y):
        return np.asarray(y, dtype=np.float32) * self.std_ + self.mean_


def resolve_existing(path: str | Path) -> Path:
    path = Path(path).expanduser()
    if path.exists():
        return path
    gz_path = Path(f"{path}.gz")
    if gz_path.exists():
        return gz_path
    if path.suffix == ".gz":
        plain_path = Path(str(path)[:-3])
        if plain_path.exists():
            return plain_path
    raise FileNotFoundError(path)


def read_csv_auto(path: str | Path, *args, **kwargs) -> pd.DataFrame:
    return pd.read_csv(resolve_existing(path), *args, **kwargs)


def write_csv_auto(df: pd.DataFrame, path: str | Path, index: bool = False) -> Path:
    path = Path(path).expanduser()
    path.parent.mkdir(parents=True, exist_ok=True)
    df.to_csv(path, index=index)
    return path


def torch_load(path: str | Path, map_location):
    try:
        return torch.load(resolve_existing(path), map_location=map_location, weights_only=False)
    except TypeError:
        return torch.load(resolve_existing(path), map_location=map_location)


def set_global_seed(seed: int) -> None:
    random.seed(seed)
    np.random.seed(seed)
    torch.manual_seed(seed)
    torch.cuda.manual_seed_all(seed)
    torch.backends.cudnn.deterministic = True
    torch.backends.cudnn.benchmark = False


def build_featurizer() -> MultipleFeaturizer:
    return MultipleFeaturizer(
        [
            cf.Stoichiometry(),
            cf.ElementProperty.from_preset("magpie"),
            cf.ValenceOrbital(props=["avg"]),
            cf.IonProperty(fast=True),
        ]
    )


def detect_formula_column(df: pd.DataFrame, explicit: str | None) -> str:
    if explicit:
        if explicit not in df.columns:
            raise ValueError(f"Formula column was not found: {explicit}")
        return explicit

    candidates = ["pretty_comp", "formula", "composition", "comp", "material", "Material"]
    for col in candidates:
        if col in df.columns:
            return col

    clean_cols = [c for c in df.columns if not str(c).startswith("Unnamed")]
    if not clean_cols:
        raise ValueError("No usable formula column was found.")
    return clean_cols[0]


def detect_target_column(df: pd.DataFrame, formula_col: str, explicit: str | None) -> str:
    if explicit:
        if explicit not in df.columns:
            raise ValueError(f"Target column was not found: {explicit}")
        return explicit

    if "property" in df.columns:
        return "property"

    numeric_cols = []
    for col in df.columns:
        if col == formula_col or str(col).startswith("Unnamed"):
            continue
        series = pd.to_numeric(df[col], errors="coerce")
        if series.notna().sum() > 0:
            numeric_cols.append(col)

    if len(numeric_cols) == 1:
        return numeric_cols[0]
    raise ValueError("Target column is ambiguous. Please pass --target_col.")


def featurize_formulas(df: pd.DataFrame, formula_col: str):
    featurizer = build_featurizer()
    labels = featurizer.feature_labels()
    features = []
    valid_indices = []
    errors = {}

    for idx, formula in df[formula_col].items():
        try:
            comp = Composition(str(formula))
            feat = featurizer.featurize(comp)
            if len(feat) != Config.input_dim:
                raise ValueError(f"expected {Config.input_dim} descriptors, got {len(feat)}")
            if any(pd.isna(x) for x in feat):
                raise ValueError("descriptor contains NaN")
            features.append(feat)
            valid_indices.append(idx)
        except Exception as exc:
            errors[int(idx)] = str(exc)

    if not valid_indices:
        raise ValueError("No valid formulas remained after descriptor generation.")

    feature_array = np.asarray(features, dtype=np.float32)
    valid_df = df.loc[valid_indices].copy()
    valid_df.insert(0, "_source_index", valid_indices)
    return valid_df.reset_index(drop=True), feature_array, labels, errors


def fit_standard_scaler(x: np.ndarray) -> dict:
    scaler = StandardScaler().fit(x)
    return {
        "mean": scaler.mean_.astype(np.float32),
        "scale": scaler.scale_.astype(np.float32),
    }


def apply_standard_scaler(x: np.ndarray, scaler_state: dict) -> np.ndarray:
    mean = np.asarray(scaler_state["mean"], dtype=np.float32)
    scale = np.asarray(scaler_state["scale"], dtype=np.float32)
    scale = np.where(np.abs(scale) < 1e-8, 1.0, scale)
    return (x.astype(np.float32) - mean) / scale


def load_pretrained_models(fusion_weights: Path, student_weights: Path, device: str):
    fusion_model = ADAE_ResNet(input_dim=Config.input_dim, latent_dim=Config.latent_dim).to(device)
    fusion_state = torch_load(fusion_weights, map_location=device)
    if isinstance(fusion_state, dict) and "model_state_dict" in fusion_state:
        fusion_state = fusion_state["model_state_dict"]
    fusion_model.load_state_dict(fusion_state)
    fusion_model.eval()

    student_model = StudentEncoder(input_dim=Config.input_dim, output_dim=Config.structure_dim).to(device)
    student_model.load_state_dict(torch_load(student_weights, map_location=device))
    student_model.eval()
    return fusion_model, student_model


@torch.no_grad()
def make_distal_features(raw_x: np.ndarray, input_scaler_state: dict, fusion_model, student_model, batch_size: int, device: str):
    scaled_x = apply_standard_scaler(raw_x, input_scaler_state)
    z_parts = []
    s_parts = []

    for start in range(0, len(raw_x), batch_size):
        stop = start + batch_size
        x_scaled_t = torch.tensor(scaled_x[start:stop], dtype=torch.float32, device=device)
        x_raw_t = torch.tensor(raw_x[start:stop], dtype=torch.float32, device=device)

        _, z_t, _ = fusion_model(x_scaled_t, epoch_ratio=1.0)
        s_t = student_model(x_raw_t)

        z_parts.append(z_t.detach().cpu().numpy().astype(np.float32))
        s_parts.append(s_t.detach().cpu().numpy().astype(np.float32))

    z = np.concatenate(z_parts, axis=0)
    s = np.concatenate(s_parts, axis=0)
    return np.concatenate([z, scaled_x.astype(np.float32), s], axis=1).astype(np.float32)


def train_one_fold(x_train, y_train, x_val, y_val, fold_id: int, args, device: str):
    set_global_seed(args.seed + fold_id)

    feature_scaler = StandardScaler().fit(x_train)
    x_tr = feature_scaler.transform(x_train).astype(np.float32)
    x_va = feature_scaler.transform(x_val).astype(np.float32)

    target_scaler = TargetScaler().fit(y_train)
    y_tr_scaled = target_scaler.transform(y_train)

    train_dataset = TensorDataset(
        torch.tensor(x_tr, dtype=torch.float32),
        torch.tensor(y_tr_scaled, dtype=torch.float32).view(-1, 1),
    )
    train_loader = DataLoader(
        train_dataset,
        batch_size=min(args.batch_size, max(16, len(train_dataset))),
        shuffle=True,
        num_workers=0,
        pin_memory=(device == "cuda"),
        drop_last=False,
    )

    model = PropertyPredictor(input_dim=x_tr.shape[1]).to(device)
    optimizer = torch.optim.AdamW(model.parameters(), lr=args.lr, weight_decay=args.weight_decay)
    scheduler = ReduceLROnPlateau(
        optimizer,
        mode="min",
        factor=0.5,
        patience=args.lr_patience,
        min_lr=args.min_lr,
    )
    criterion = nn.SmoothL1Loss(beta=0.5)
    x_val_t = torch.tensor(x_va, dtype=torch.float32, device=device)

    best_val_mae = float("inf")
    best_state = None
    no_improve = 0

    for _ in range(args.max_epochs):
        model.train()
        for xb, yb in train_loader:
            xb = xb.to(device, non_blocking=True)
            yb = yb.to(device, non_blocking=True)

            optimizer.zero_grad(set_to_none=True)
            loss = criterion(model(xb), yb)
            loss.backward()
            torch.nn.utils.clip_grad_norm_(model.parameters(), args.grad_clip)
            optimizer.step()

        model.eval()
        with torch.no_grad():
            val_pred = model(x_val_t).detach().cpu().numpy().reshape(-1)
            val_pred = target_scaler.inverse_transform(val_pred)
            val_mae = mean_absolute_error(y_val, val_pred)

        scheduler.step(val_mae)
        if val_mae < best_val_mae - Config.min_delta:
            best_val_mae = float(val_mae)
            best_state = copy.deepcopy(model.state_dict())
            no_improve = 0
        else:
            no_improve += 1
            if no_improve >= args.patience:
                break

    if best_state is None:
        best_state = copy.deepcopy(model.state_dict())

    return {
        "fold_id": int(fold_id),
        "state_dict": best_state,
        "best_val_mae": float(best_val_mae),
        "feature_scaler_mean": feature_scaler.mean_.astype(np.float32),
        "feature_scaler_scale": feature_scaler.scale_.astype(np.float32),
        "target_scaler_mean": float(target_scaler.mean_),
        "target_scaler_std": float(target_scaler.std_),
        "input_dim": int(x_tr.shape[1]),
    }


def train_ensemble(features: np.ndarray, target: np.ndarray, args, device: str):
    if len(features) < 3:
        raise ValueError("At least three training rows are required.")

    folds = min(args.folds, len(features))
    if folds < 2:
        folds = 2

    kf = KFold(n_splits=folds, shuffle=True, random_state=args.seed)
    fold_bundles = []

    for fold_id, (train_idx, val_idx) in enumerate(kf.split(features), start=1):
        fold_bundle = train_one_fold(
            features[train_idx],
            target[train_idx],
            features[val_idx],
            target[val_idx],
            fold_id,
            args,
            device,
        )
        fold_bundles.append(fold_bundle)
        print(f"Fold {fold_id:02d}/{folds} | validation MAE = {fold_bundle['best_val_mae']:.6g}")

    return fold_bundles


@torch.no_grad()
def predict_ensemble(feature_matrix: np.ndarray, fold_bundles: list[dict], device: str):
    predictions = []
    for fold in fold_bundles:
        model = PropertyPredictor(input_dim=int(fold["input_dim"])).to(device)
        model.load_state_dict(fold["state_dict"])
        model.eval()

        scaler_state = {
            "mean": fold["feature_scaler_mean"],
            "scale": fold["feature_scaler_scale"],
        }
        x = apply_standard_scaler(feature_matrix, scaler_state)
        x_t = torch.tensor(x, dtype=torch.float32, device=device)
        pred = model(x_t).detach().cpu().numpy().reshape(-1)
        pred = pred * float(fold["target_scaler_std"]) + float(fold["target_scaler_mean"])
        predictions.append(pred.astype(np.float32))

    pred_stack = np.stack(predictions, axis=0)
    return pred_stack.mean(axis=0), pred_stack.std(axis=0)


def append_invalid_rows(result_df: pd.DataFrame, source_df: pd.DataFrame, errors: dict[int, str]) -> pd.DataFrame:
    if not errors:
        return result_df

    invalid = source_df.loc[list(errors.keys())].copy()
    invalid.insert(0, "_source_index", list(errors.keys()))
    invalid["prediction"] = np.nan
    invalid["ensemble_std"] = np.nan
    invalid["parse_error"] = invalid["_source_index"].map(lambda i: errors[int(i)])
    merged = pd.concat([result_df, invalid], ignore_index=True, sort=False)
    return merged.sort_values("_source_index").reset_index(drop=True)


def run(args) -> None:
    device = args.device if args.device == "cpu" or torch.cuda.is_available() else "cpu"
    set_global_seed(args.seed)

    train_df = read_csv_auto(args.train_csv)
    predict_df = read_csv_auto(args.predict_csv)
    train_formula_col = detect_formula_column(train_df, args.train_formula_col or args.formula_col)
    predict_formula_col = detect_formula_column(predict_df, args.predict_formula_col or args.formula_col)
    target_col = detect_target_column(train_df, train_formula_col, args.target_col)

    valid_train_df, raw_x_train, descriptor_labels, train_errors = featurize_formulas(train_df, train_formula_col)
    y_train = pd.to_numeric(valid_train_df[target_col], errors="coerce").to_numpy(dtype=np.float32)
    valid_target_mask = np.isfinite(y_train)
    valid_train_df = valid_train_df.loc[valid_target_mask].reset_index(drop=True)
    raw_x_train = raw_x_train[valid_target_mask]
    y_train = y_train[valid_target_mask]
    if len(y_train) < 3:
        raise ValueError("At least three valid labeled training rows are required.")

    valid_predict_df, raw_x_predict, _, predict_errors = featurize_formulas(predict_df, predict_formula_col)

    input_scaler_state = fit_standard_scaler(raw_x_train)
    fusion_model, student_model = load_pretrained_models(args.fusion_weights, args.student_weights, device)

    train_features = make_distal_features(
        raw_x_train,
        input_scaler_state,
        fusion_model,
        student_model,
        args.feature_batch_size,
        device,
    )
    predict_features = make_distal_features(
        raw_x_predict,
        input_scaler_state,
        fusion_model,
        student_model,
        args.feature_batch_size,
        device,
    )

    fold_bundles = train_ensemble(train_features, y_train, args, device)
    mean_pred, std_pred = predict_ensemble(predict_features, fold_bundles, device)

    result_df = valid_predict_df.copy()
    result_df["prediction"] = mean_pred
    result_df["ensemble_std"] = std_pred

    if args.predict_target_col and args.predict_target_col in result_df.columns:
        y_true = pd.to_numeric(result_df[args.predict_target_col], errors="coerce").to_numpy(dtype=np.float32)
        if np.isfinite(y_true).all():
            result_df["absolute_error"] = np.abs(y_true - result_df["prediction"].to_numpy(dtype=np.float32))
            print(f"Prediction MAE = {mean_absolute_error(y_true, result_df['prediction']):.6g}")

    result_df = append_invalid_rows(result_df, predict_df, predict_errors)
    output_csv = write_csv_auto(result_df, args.output_csv)

    bundle = {
        "feature_mode": "z+x+s",
        "descriptor_labels": descriptor_labels,
        "input_scaler_mean": input_scaler_state["mean"],
        "input_scaler_scale": input_scaler_state["scale"],
        "folds": fold_bundles,
        "train_formula_col": train_formula_col,
        "predict_formula_col": predict_formula_col,
        "target_col": target_col,
        "train_rows_used": int(len(y_train)),
        "train_rows_skipped": int(len(train_errors)),
    }
    args.bundle.parent.mkdir(parents=True, exist_ok=True)
    torch.save(bundle, args.bundle)

    print(f"Training rows used: {len(y_train)}")
    print(f"Training rows skipped by formula parsing: {len(train_errors)}")
    print(f"Candidate rows predicted: {len(valid_predict_df)}")
    print(f"Candidate rows skipped by formula parsing: {len(predict_errors)}")
    print(f"DISTAL feature dimension: {train_features.shape[1]}")
    print(f"Saved predictions to: {output_csv}")
    print(f"Saved task bundle to: {args.bundle}")


def parse_args():
    parser = argparse.ArgumentParser(
        description=(
            "Train a task-specific DISTAL property predictor from a labeled CSV "
            "and predict a second CSV of candidate formulas."
        )
    )
    parser.add_argument("--train_csv", type=Path, required=True, help="CSV with known formulas and target property values.")
    parser.add_argument("--predict_csv", type=Path, required=True, help="CSV with formulas to predict.")
    parser.add_argument("--formula_col", type=str, default=None, help="Formula column used by both CSV files.")
    parser.add_argument("--train_formula_col", type=str, default=None, help="Formula column in --train_csv.")
    parser.add_argument("--predict_formula_col", type=str, default=None, help="Formula column in --predict_csv.")
    parser.add_argument("--target_col", type=str, default=None, help="Target property column in --train_csv.")
    parser.add_argument("--predict_target_col", type=str, default=None, help="Optional true target column in --predict_csv.")
    parser.add_argument("--fusion_weights", type=Path, default=DEFAULT_FUSION_WEIGHTS)
    parser.add_argument("--student_weights", type=Path, default=DEFAULT_STUDENT_WEIGHTS)
    parser.add_argument("--output_csv", type=Path, default=DEFAULT_OUTPUT_CSV)
    parser.add_argument("--bundle", type=Path, default=DEFAULT_BUNDLE_PATH)
    parser.add_argument("--device", type=str, default=DEVICE, choices=["cpu", "cuda"])
    parser.add_argument("--folds", type=int, default=10)
    parser.add_argument("--seed", type=int, default=2024)
    parser.add_argument("--batch_size", type=int, default=128)
    parser.add_argument("--feature_batch_size", type=int, default=4096)
    parser.add_argument("--max_epochs", type=int, default=700)
    parser.add_argument("--patience", type=int, default=80)
    parser.add_argument("--lr_patience", type=int, default=25)
    parser.add_argument("--lr", type=float, default=3e-4)
    parser.add_argument("--min_lr", type=float, default=1e-6)
    parser.add_argument("--weight_decay", type=float, default=5e-6)
    parser.add_argument("--grad_clip", type=float, default=1.0)
    return parser.parse_args()


if __name__ == "__main__":
    run(parse_args())
