import numpy as np

# -------------------------------------------------------------------
# 1. Put your digitized data here
#    Units:
#      - B in Tesla
#      - frequency in kHz (x-axis of your plot)
#      - Pv in kW/m^3 (y-axis of your plot)
#
# IMPORTANT:
#   The numbers below are ONLY EXAMPLES / PLACEHOLDERS.
#   Replace each "f_khz" and "Pv_kW_m3" array with points
#   read from your datasheet plot (solid 25 °C lines).
# -------------------------------------------------------------------

core_loss_data_25C = {
    0.025: {  # 25 mT
        "f_khz": np.array([2.93084374470676,
            8.7361542468832,
            44.706096185109,
            125.963436020903,
            193.225112880907,
            ], dtype=float),
        "Pv_kW_m3": np.array([0.177193262696068,
            0.679670055798598,
            5.03342212801684,
            17.4677120801992,
            30.0788251804309,
            ], dtype=float),
    },
    0.050: {  # 50 mT
        "f_khz": np.array([2.9975926563905,
            9.24208952674728,
            38.618648046052,
            188.92247157063,
            ], dtype=float),
        "Pv_kW_m3": np.array([0.917768017464649,
            3.47034515205292,
            19.0328184767803,
            123.927503749051,
            ], dtype=float),
    },
    0.100: {  # 100 mT
        "f_khz": np.array([2.48,
            8.84,
            34.12,
            87.85,
            193.23,
            ], dtype=float),
        "Pv_kW_m3": np.array([3.14,
            13.89,
            69.94,
            216.47,
            564.36,
            ], dtype=float),
    },
    0.200: {  # 200 mT
        "f_khz": np.array([2.14,
                8.64,
                32.25,
                83.05,
                195.41,
                ], dtype=float),
        "Pv_kW_m3": np.array([10.44,
            56.44,
            276.05,
            891.89,
            2497.56,
            ], dtype=float),
    },
    0.300: {  # 300 mT
        "f_khz": np.array([2.19,
            5.03,
            12.39,
            24.89,
            49.47,
            ], dtype=float),
        "Pv_kW_m3": np.array([26.83,
            67.97,
            193.07,
            442.55,
            1000,
            ], dtype=float),
    },
}


# -------------------------------------------------------------------
# 2. Helper: 1-D log-log interpolation in frequency
# -------------------------------------------------------------------

def _loglog_interp(x, xp, fp):
    """
    Log-log interpolation: fp vs xp, evaluated at x.

    All arguments are assumed > 0.
    """
    x = np.asarray(x, dtype=float)
    xp = np.asarray(xp, dtype=float)
    fp = np.asarray(fp, dtype=float)

    logx = np.log10(x)
    logxp = np.log10(xp)
    logfp = np.log10(fp)

    logy = np.interp(logx, logxp, logfp)
    return 10.0 ** logy


# -------------------------------------------------------------------
# 3. Main function: Pv(B, f) interpolation
# -------------------------------------------------------------------

def core_loss_Pv(B_T, f_Hz, data=core_loss_data_25C):
    """
    Interpolate core loss density P_v for given flux density and frequency.

    Parameters
    ----------
    B_T : float
        Peak flux density in Tesla (e.g. 0.05 for 50 mT).
    f_Hz : float or array_like
        Frequency in Hz.
    data : dict
        Dict mapping B_T -> {"f_khz": array, "Pv_kW_m3": array}.

    Returns
    -------
    Pv_kW_m3 : float or ndarray
        Core loss density in kW/m^3 at 25 °C.
    """
    B_T = float(B_T)
    f_khz = np.asarray(f_Hz, dtype=float) / 1e3  # plot uses kHz

    if np.any(f_khz <= 0):
        raise ValueError("Frequency must be > 0 Hz")

    # Sorted list of available B values from your digitized curves
    B_list = np.array(sorted(data.keys()))

    # Handle out-of-range B values
    if B_T < B_list.min():
        raise ValueError(
            f"B={B_T:.3f} T is below available range "
            f"[{B_list.min():.3f}, {B_list.max():.3f}] T"
        )

    # For B > max data (0.3T), extrapolate up to 0.4T using two highest curves
    # If B > 0.4T, clamp to the extrapolated value at 0.4T
    if B_T > B_list.max():
        # Use the two highest B values for extrapolation
        B1 = B_list[-2]  # 0.2T
        B2 = B_list[-1]  # 0.3T

        d1 = data[B1]
        d2 = data[B2]

        # Interpolate along frequency (log-log) on each B curve
        Pv1 = _loglog_interp(f_khz, d1["f_khz"], d1["Pv_kW_m3"])
        Pv2 = _loglog_interp(f_khz, d2["f_khz"], d2["Pv_kW_m3"])

        # Extrapolate in log(Pv) vs B (Pv ~ B^b approximately)
        logPv1 = np.log10(Pv1)
        logPv2 = np.log10(Pv2)

        # Clamp B_T to 0.4T for extrapolation calculation
        B_T_clamped = min(B_T, 0.4)
        alpha = (B_T_clamped - B1) / (B2 - B1)
        logPv = (1.0 - alpha) * logPv1 + alpha * logPv2

        return 10.0 ** logPv

    # If B_T is (approximately) one of the tabulated curves, just interpolate in f
    idx_match = np.where(np.isclose(B_list, B_T, rtol=1e-4, atol=1e-6))[0]
    if idx_match.size:
        B_key = B_list[idx_match[0]]
        d = data[B_key]
        return _loglog_interp(f_khz, d["f_khz"], d["Pv_kW_m3"])

    # Otherwise, find the two B curves that bracket B_T
    idx = np.searchsorted(B_list, B_T)
    B1 = B_list[idx - 1]
    B2 = B_list[idx]

    d1 = data[B1]
    d2 = data[B2]

    # Interpolate along frequency (log-log) on each B curve
    Pv1 = _loglog_interp(f_khz, d1["f_khz"], d1["Pv_kW_m3"])
    Pv2 = _loglog_interp(f_khz, d2["f_khz"], d2["Pv_kW_m3"])

    # Now interpolate between B1 and B2.
    # Do it in log(Pv) vs B (Pv ~ B^b approximately).
    logPv1 = np.log10(Pv1)
    logPv2 = np.log10(Pv2)

    alpha = (B_T - B1) / (B2 - B1)
    logPv = (1.0 - alpha) * logPv1 + alpha * logPv2

    return 10.0 ** logPv



if __name__ == '__main__':
    p = core_loss_Pv(0.2, 25_000)
    print (p)