from __future__ import annotations

import numpy as np
import pandas as pd

from .utils import infer_family


def sample_base_designs(n=1000, seed=7) -> pd.DataFrame:
    rng = np.random.default_rng(seed)
    types = np.array([
        "NC_pillbox_Lband",
        "SC_elliptical_1p3GHz",
        "NC_Cband_TW",
        "NC_Xband_SW",
    ])
    probs = np.array([0.25, 0.20, 0.30, 0.25])
    types_arr = rng.choice(types, size=n, p=probs)

    f = np.zeros(n)
    rq = np.zeros(n)
    q0 = np.zeros(n)
    eacc = np.zeros(n)
    length = np.zeros(n)
    beta = np.zeros(n)
    cryo = np.full(n, np.nan)

    for i, typ in enumerate(types_arr):
        if typ == "NC_pillbox_Lband":
            f[i] = rng.uniform(0.20e9, 0.80e9)
            rq[i] = rng.uniform(60, 150)
            q0[i] = rng.uniform(1e4, 8e4)
            eacc[i] = rng.uniform(8, 20)
            length[i] = rng.uniform(0.15, 0.60)
            beta[i] = rng.uniform(0.7, 2.5)
        elif typ == "SC_elliptical_1p3GHz":
            f[i] = rng.uniform(0.80e9, 1.60e9)
            rq[i] = rng.uniform(200, 400)
            q0[i] = rng.uniform(5e9, 2e10)
            eacc[i] = rng.uniform(10, 25)
            length[i] = rng.uniform(0.60, 1.20)
            beta[i] = rng.uniform(0.8, 2.0)
            cryo[i] = rng.uniform(1.8, 4.2)
        elif typ == "NC_Cband_TW":
            f[i] = rng.uniform(5.00e9, 6.00e9)
            rq[i] = rng.uniform(100, 280)
            q0[i] = rng.uniform(2e4, 8e4)
            eacc[i] = rng.uniform(12, 30)
            length[i] = rng.uniform(0.20, 0.60)
            beta[i] = rng.uniform(0.8, 3.0)
        else:
            f[i] = rng.uniform(9.00e9, 12.0e9)
            rq[i] = rng.uniform(80, 200)
            q0[i] = rng.uniform(1e4, 6e4)
            eacc[i] = rng.uniform(20, 45)
            length[i] = rng.uniform(0.15, 0.45)
            beta[i] = rng.uniform(0.8, 3.0)

    df = pd.DataFrame({
        "type": types_arr,
        "freq_Hz": f,
        "freq_GHz": f * 1e-9,
        "R_over_Q_ohm": rq,
        "Q0": q0,
        "Eacc_MVpm": eacc,
        "L_m": length,
        "beta": beta,
        "cryo_temp_K": cryo,
    })
    df["family"] = df["type"].map(infer_family)
    df["Vc_MV"] = df["Eacc_MVpm"] * df["L_m"]
    df["Q_loaded"] = df["Q0"] / (1.0 + df["beta"])
    return df


def sample_family_conditioned_operating_point(df: pd.DataFrame, seed=17) -> pd.DataFrame:
    rng = np.random.default_rng(seed)

    pulse_ns = np.zeros(len(df))
    rep_hz = np.zeros(len(df))
    p_aux_kW = np.zeros(len(df))
    source_power_avail_kW = np.zeros(len(df))
    cooling_capacity_kW = np.zeros(len(df))
    surf_factor = np.zeros(len(df))
    geom_factor = np.zeros(len(df))
    freq_factor = np.zeros(len(df))
    fab_sigma_um = np.zeros(len(df))
    delta_allow_um = np.zeros(len(df))
    s_f = np.zeros(len(df))
    s_phi = np.zeros(len(df))
    s_c = np.zeros(len(df))

    for i, fam in enumerate(df["family"].values):
        if fam == "Lband":
            pulse_ns[i] = rng.uniform(800.0, 3000.0)
            rep_hz[i] = rng.uniform(100.0, 2000.0)
            p_aux_kW[i] = rng.uniform(0.4, 1.5)
            source_power_avail_kW[i] = rng.uniform(4.0, 8.0)
            cooling_capacity_kW[i] = rng.uniform(6.0, 12.0)
            surf_factor[i] = rng.uniform(0.9, 1.1)
            geom_factor[i] = rng.uniform(0.8, 1.1)
            freq_factor[i] = rng.uniform(0.9, 1.05)
            fab_sigma_um[i] = rng.uniform(10.0, 40.0)
            delta_allow_um[i] = rng.uniform(40.0, 100.0)
            s_f[i] = rng.uniform(0.003, 0.010)
            s_phi[i] = rng.uniform(0.002, 0.008)
            s_c[i] = rng.uniform(0.001, 0.006)
        elif fam == "Cband":
            pulse_ns[i] = rng.uniform(200.0, 800.0)
            rep_hz[i] = rng.uniform(50.0, 500.0)
            p_aux_kW[i] = rng.uniform(0.8, 2.2)
            source_power_avail_kW[i] = rng.uniform(9.0, 15.0)
            cooling_capacity_kW[i] = rng.uniform(10.0, 20.0)
            surf_factor[i] = rng.uniform(0.95, 1.20)
            geom_factor[i] = rng.uniform(0.9, 1.2)
            freq_factor[i] = rng.uniform(1.0, 1.15)
            fab_sigma_um[i] = rng.uniform(5.0, 20.0)
            delta_allow_um[i] = rng.uniform(15.0, 40.0)
            s_f[i] = rng.uniform(0.010, 0.025)
            s_phi[i] = rng.uniform(0.008, 0.020)
            s_c[i] = rng.uniform(0.006, 0.015)
        elif fam == "Xband":
            pulse_ns[i] = rng.uniform(50.0, 300.0)
            rep_hz[i] = rng.uniform(10.0, 200.0)
            p_aux_kW[i] = rng.uniform(1.0, 3.0)
            source_power_avail_kW[i] = rng.uniform(12.0, 20.0)
            cooling_capacity_kW[i] = rng.uniform(12.0, 24.0)
            surf_factor[i] = rng.uniform(1.0, 1.35)
            geom_factor[i] = rng.uniform(1.0, 1.3)
            freq_factor[i] = rng.uniform(1.1, 1.3)
            fab_sigma_um[i] = rng.uniform(2.0, 10.0)
            delta_allow_um[i] = rng.uniform(5.0, 18.0)
            s_f[i] = rng.uniform(0.020, 0.060)
            s_phi[i] = rng.uniform(0.015, 0.050)
            s_c[i] = rng.uniform(0.010, 0.035)
        elif fam == "SC":
            pulse_ns[i] = rng.uniform(5000.0, 100000.0)
            rep_hz[i] = rng.uniform(1.0, 100.0)
            p_aux_kW[i] = rng.uniform(4.0, 15.0)
            source_power_avail_kW[i] = rng.uniform(8.0, 20.0)
            cooling_capacity_kW[i] = rng.uniform(15.0, 40.0)
            surf_factor[i] = rng.uniform(0.9, 1.1)
            geom_factor[i] = rng.uniform(0.8, 1.0)
            freq_factor[i] = rng.uniform(0.9, 1.05)
            fab_sigma_um[i] = rng.uniform(10.0, 30.0)
            delta_allow_um[i] = rng.uniform(30.0, 80.0)
            s_f[i] = rng.uniform(0.004, 0.012)
            s_phi[i] = rng.uniform(0.003, 0.010)
            s_c[i] = rng.uniform(0.002, 0.008)
        else:
            pulse_ns[i] = rng.uniform(100.0, 1000.0)
            rep_hz[i] = rng.uniform(10.0, 500.0)
            p_aux_kW[i] = rng.uniform(1.0, 3.0)
            source_power_avail_kW[i] = rng.uniform(5.0, 15.0)
            cooling_capacity_kW[i] = rng.uniform(8.0, 16.0)
            surf_factor[i] = 1.0
            geom_factor[i] = 1.0
            freq_factor[i] = 1.0
            fab_sigma_um[i] = 10.0
            delta_allow_um[i] = 20.0
            s_f[i] = 0.01
            s_phi[i] = 0.01
            s_c[i] = 0.01

    out = df.copy()
    out["pulse_length_ns"] = pulse_ns
    out["rep_rate_Hz"] = rep_hz
    out["P_aux_kW"] = p_aux_kW
    out["source_power_avail_kW"] = source_power_avail_kW
    out["cooling_capacity_kW"] = cooling_capacity_kW
    out["surf_factor"] =  surf_factor
    out["geom_factor"] = geom_factor
    out["freq_factor"] = freq_factor
    out["fab_sigma_um"] = fab_sigma_um
    out["delta_allow_um"] = delta_allow_um
    out["S_f"] = s_f
    out["S_phi"] = s_phi
    out["S_c"] = s_c
    return out