import numpy as np
import pandas as pd

def cfg_cylindrical():
    return {
        "cols": ["diameter", "height", "fissile_concentration", "isotopic_abundance"],
        "label": "keff",
        "aliases": {
            "直径": "diameter",
            "高度": "height",
            "铀浓度": "fissile_concentration",
            "U浓度": "fissile_concentration",
            "铀富集度": "isotopic_abundance",
            "U富集度": "isotopic_abundance",
            # 兼容旧英文键
            "u_concentration": "fissile_concentration",
            "u_enrichment": "isotopic_abundance"
        }
    }

def cfg_ring():
    return {
        "cols": ["diameter", "height", "fissile_concentration", "isotopic_abundance"],
        "label": "keff",
        "aliases": {
            "外径": "diameter",
            "高度": "height",
            "Pu浓度": "fissile_concentration",
            "Pu240占比": "isotopic_abundance",
            # 兼容旧英文键
            "pu_concentration": "fissile_concentration",
            "pu_isotope": "isotopic_abundance"
        }
    }

TYPE_CONFIG = {
    "CylindricalTank": cfg_cylindrical(),
    "AnnularTank": cfg_ring()
}

def feature_names(device_type):
    cfg = TYPE_CONFIG[device_type]
    base = cfg["cols"]
    if device_type in ("CylindricalTank", "AnnularTank"):
        derived = [f"{base[0]}*{base[1]}", f"{base[2]}*{base[3]}"]
        return base + derived
    return base

def derive_features(device_type, X):
    if device_type in ("CylindricalTank", "AnnularTank"):
        return np.hstack([X, (X[:, 0] * X[:, 1]).reshape(-1, 1), (X[:, 2] * X[:, 3]).reshape(-1, 1)])
    return X

def _resolve_col(df, key, aliases):
    if key in df.columns:
        return df[key].values
    for a, k in aliases.items():
        if k == key and a in df.columns:
            return df[a].values
    raise KeyError(key)

def load_from_excel(device_type, path):
    cfg = TYPE_CONFIG[device_type]
    df = pd.read_excel(path)
    cols = cfg["cols"]
    aliases = cfg.get("aliases", {})
    arrs = []
    for k in cols:
        v = _resolve_col(df, k, aliases)
        arrs.append(v.astype(float))
    X = np.vstack(arrs).T
    y = df[cfg["label"]].values.astype(float) if cfg["label"] in df.columns else df[cfg["label"]].values.astype(float)
    Xd = derive_features(device_type, X)
    return Xd, y

def features_to_vector(device_type, features):
    cfg = TYPE_CONFIG[device_type]
    aliases = cfg.get("aliases", {})
    vals = []
    for k in cfg["cols"]:
        v = features.get(k)
        if v is None:
            for a, ck in aliases.items():
                if ck == k and a in features:
                    v = features[a]
                    break
        vals.append(float(v))
    X = np.array(vals, dtype=float).reshape(1, -1)
    Xd = derive_features(device_type, X)
    return Xd