import itertools
import math
import os
import random
from dataclasses import dataclass
from typing import Dict, List, Sequence, Tuple

import numpy as np
from matplotlib import pyplot as plt
from scipy.interpolate import CubicSpline

plt.rcParams["font.sans-serif"] = ["DejaVu Sans", "Arial", "Helvetica"]
plt.rcParams["axes.unicode_minus"] = False


# -----------------------------
# 你的PDMS光学常数模型
# -----------------------------
def make_strictly_increasing(wl, n, k):
    # 去除完全相同的点
    unique_wl, indices = np.unique(wl, return_index=True)
    if len(unique_wl) != len(wl):
        print(f"Removed {len(wl) - len(unique_wl)} duplicate wavelength points")
        wl, n, k = wl[indices], n[indices], k[indices]

    # 确保严格递增
    is_increasing = np.diff(wl) > 0
    if not all(is_increasing):
        # 移除不满足严格递增的点
        valid_indices = np.concatenate([[True], is_increasing])
        wl, n, k = wl[valid_indices], n[valid_indices], k[valid_indices]

    return wl, n, k


def read_split_data(file_path):
    """
    解析分块数据格式：
    第一部分：wl	n（多行数据）
    第二部分：wl	k（多行数据）
    返回：sorted_wl, n, k（波长已排序，保证n和k一一对应）
    """
    with open(file_path, 'r', encoding='utf-8') as f:
        lines = [line.strip() for line in f if line.strip() and not line.startswith('#')]

    # 分割n数据和k数据（以"wl	k"为分界）
    split_idx = None
    for i, line in enumerate(lines):
        if line == 'wl	k':  # 找到k数据的表头
            split_idx = i
            break

    # 提取n数据（表头后到split_idx前）
    n_lines = lines[1:split_idx]  # 跳过"wl	n"表头
    wl_n = []
    n_list = []
    for line in n_lines:
        wl, n_val = line.split()
        wl_n.append(float(wl))
        n_list.append(float(n_val))

    # 提取k数据（split_idx后）
    k_lines = lines[split_idx + 1:]  # 跳过"wl	k"表头
    wl_k = []
    k_list = []
    for line in k_lines:
        wl, k_val = line.split()
        wl_k.append(float(wl))
        k_list.append(float(k_val))

    # 转换为numpy数组
    wl_n = np.array(wl_n)
    n_list = np.array(n_list)
    wl_k = np.array(wl_k)
    k_list = np.array(k_list)

    # 确保n和k的波长完全一致（否则插值会出错）
    assert np.allclose(wl_n, wl_k), "n和k的波长列表不一致！"

    # 排序（按波长递增，避免插值异常）
    sorted_idx = np.argsort(wl_n)
    sorted_wl = wl_n[sorted_idx]
    sorted_n = n_list[sorted_idx]
    sorted_k = k_list[sorted_idx]

    return sorted_wl, sorted_n, sorted_k


# 全局变量存储PDMS插值函数
_PDMS_CS_N = None
_PDMS_CS_K = None
_PDMS_WL_MIN = None
_PDMS_WL_MAX = None


def initialize_pdms_model(data_file_path):
    """
    初始化PDMS光学常数模型
    """
    global _PDMS_CS_N, _PDMS_CS_K, _PDMS_WL_MIN, _PDMS_WL_MAX

    # 读取数据
    wl_all, n_all, k_all = read_split_data(data_file_path)
    wl_all, n_all, k_all = make_strictly_increasing(wl_all, n_all, k_all)

    # 创建插值函数
    _PDMS_CS_N = CubicSpline(wl_all, n_all)
    _PDMS_CS_K = CubicSpline(wl_all, k_all)
    _PDMS_WL_MIN = wl_all.min()
    _PDMS_WL_MAX = wl_all.max()

    print(f"PDMS模型初始化完成: 波长范围 [{_PDMS_WL_MIN:.2f}, {_PDMS_WL_MAX:.2f}] μm")


def pdms_nk(wavelength_um: np.ndarray) -> np.ndarray:
    """
    使用你的PDMS模型计算复折射率
    """
    global _PDMS_CS_N, _PDMS_CS_K, _PDMS_WL_MIN, _PDMS_WL_MAX

    if _PDMS_CS_N is None:
        raise ValueError("PDMS模型未初始化，请先调用initialize_pdms_model()")

    # 确保波长在有效范围内
    wl_clipped = np.clip(wavelength_um, _PDMS_WL_MIN, _PDMS_WL_MAX)

    # 计算n和k
    n = _PDMS_CS_N(wl_clipped)
    k = _PDMS_CS_K(wl_clipped)

    # 返回复折射率
    return n - 1j * k


# 初始化PDMS模型（请修改为你的实际数据文件路径）
initialize_pdms_model('/Users/spasolreisa/IdeaProjects/asiaMath/data.txt')


def planck_weight(wavelength_um: np.ndarray, temperature: float = 300.0) -> np.ndarray:
    wl_m = wavelength_um * 1e-6
    c1 = 3.7418e-16
    c2 = 1.4388e-2
    spectral = c1 / (wl_m ** 5 * (np.exp(c2 / (wl_m * temperature)) - 1))
    return spectral


def solar_weight(wavelength_um: np.ndarray) -> np.ndarray:
    center1, width1 = 0.6, 0.35
    center2, width2 = 1.6, 0.45
    return np.exp(-((wavelength_um - center1) / width1) ** 2) + 0.35 * np.exp(
        -((wavelength_um - center2) / width2) ** 2
    )


@dataclass
class Material:
    name: str
    n_const: float
    k_const: float

    def nk(self, wavelength_um: np.ndarray) -> np.ndarray:
        n = np.full_like(wavelength_um, self.n_const, dtype=np.complex128)
        k = np.full_like(wavelength_um, self.k_const, dtype=np.complex128)
        return n - 1j * k


def pdms_index(wavelength_um: np.ndarray) -> np.ndarray:
    """
    修改为使用你的PDMS模型
    """
    return pdms_nk(wavelength_um)


def ag_index(wavelength_um: np.ndarray) -> np.ndarray:
    n = 0.15 + 0.6 * np.exp(-((wavelength_um - 0.5) / 0.4) ** 2)
    k = 4.5 + 3.5 * np.exp(-((wavelength_um - 10) / 6) ** 2)
    return n - 1j * k


def transfer_matrix_stack(
        wavelength_um: np.ndarray,
        layer_nk: Sequence[np.ndarray],
        thickness_um: Sequence[float],
        substrate_nk: np.ndarray,
        n0: float = 1.0,
) -> Tuple[np.ndarray, np.ndarray, np.ndarray]:
    beta = 2 * np.pi / (wavelength_um * 1e-6)
    q0 = n0
    qs = substrate_nk

    R = np.zeros_like(wavelength_um)
    T = np.zeros_like(wavelength_um)

    for idx, wl in enumerate(wavelength_um):
        M = np.identity(2, dtype=complex)
        for nk, d in zip(layer_nk, thickness_um):
            n_layer = nk[idx]
            delta = beta[idx] * n_layer * d * 1e-6
            cos = np.cos(delta)
            sin = 1j * np.sin(delta)
            q = n_layer
            Mj = np.array([[cos, sin / q], [q * sin, cos]], dtype=complex)
            M = M @ Mj

        numerator = (
                q0 * M[0, 0]
                + q0 * qs[idx] * M[0, 1]
                - M[1, 0]
                - qs[idx] * M[1, 1]
        )
        denominator = (
                q0 * M[0, 0]
                + q0 * qs[idx] * M[0, 1]
                + M[1, 0]
                + qs[idx] * M[1, 1]
        )
        r = numerator / denominator
        t = 2 * q0 / denominator
        R[idx] = np.abs(r) ** 2
        T[idx] = np.real(qs[idx] / q0) * np.abs(t) ** 2

    A = np.clip(1 - R - T, 0, 1)
    return R, T, A


def evaluate_stack(design: Dict) -> Dict:
    solar_wl = np.linspace(0.35, 2.5, 120)
    ir_wl = np.linspace(8, 13, 200)
    solar_w = solar_weight(solar_wl)
    ir_w = planck_weight(ir_wl)

    substrate = ag_index

    layer_funcs = []
    thickness = []
    for layer in design["layers"]:
        material = layer["material"]
        thickness.append(layer["thickness"])
        if material == "PDMS":
            layer_funcs.append(pdms_index)
        else:
            mat = MATERIAL_LIBRARY[material]
            layer_funcs.append(lambda wl, m=mat: m.nk(wl))

    solar_nk = [func(solar_wl) for func in layer_funcs]
    ir_nk = [func(ir_wl) for func in layer_funcs]

    solar_R, _, solar_A = transfer_matrix_stack(
        solar_wl, solar_nk, thickness, substrate(solar_wl)
    )
    ir_R, _, ir_A = transfer_matrix_stack(ir_wl, ir_nk, thickness, substrate(ir_wl))

    alpha = float(np.trapz(solar_A * solar_w, solar_wl) / np.trapz(solar_w, solar_wl))
    epsilon = float(np.trapz(ir_A * ir_w, ir_wl) / np.trapz(ir_w, ir_wl))

    score = epsilon - 0.3 * alpha
    return {"alpha": alpha, "epsilon": epsilon, "score": score}


MATERIAL_LIBRARY: Dict[str, Material] = {
    "SiO2": Material("SiO2", 1.45, 1e-4),
    "Al2O3": Material("Al2O3", 1.76, 1.5e-3),
    "TiO2": Material("TiO2", 2.40, 5e-3),
    "Si3N4": Material("Si3N4", 2.05, 2e-3),
    "HfO2": Material("HfO2", 1.9, 2e-3),
}


def random_design() -> Dict:
    num_layers = random.choice([2, 3])
    middle_materials = random.sample(list(MATERIAL_LIBRARY.keys()), num_layers)
    layers = [{"material": "PDMS", "thickness": random.uniform(10, 50)}]
    for mat in middle_materials:
        layers.append(
            {
                "material": mat,
                "thickness": random.uniform(0.05, 2.0),
            }
        )
    return {"layers": layers}


def optimize(iterations: int = 800) -> List[Dict]:
    best_designs: List[Dict] = []
    for _ in range(iterations):
        design = random_design()
        metrics = evaluate_stack(design)
        design.update(metrics)
        best_designs.append(design)

    best_designs.sort(key=lambda x: x["score"], reverse=True)
    return best_designs[:15]


def write_summary(designs: List[Dict], path: str) -> None:
    with open(path, "w", encoding="utf-8") as f:
        f.write("rank,score,epsilon,alpha,layers\n")
        for idx, design in enumerate(designs, start=1):
            layer_desc = ";".join(
                f"{layer['material']}@{layer['thickness']:.3f}um"
                for layer in design["layers"]
            )
            f.write(
                f"{idx},{design['score']:.4f},{design['epsilon']:.4f},"
                f"{design['alpha']:.4f},{layer_desc}\n"
            )


def plot_pareto(designs: List[Dict], path: str) -> None:
    eps = [d["epsilon"] for d in designs]
    alpha = [d["alpha"] for d in designs]
    scores = [d["score"] for d in designs]
    fig, ax = plt.subplots(figsize=(6, 5))
    scatter = ax.scatter(alpha, eps, c=scores, cmap="viridis", s=80)
    ax.set_xlabel("Solar-weighted Absorption α")
    ax.set_ylabel("8-13 µm Emissivity ε")
    ax.set_title("Multilayer Design Performance Distribution")
    plt.colorbar(scatter, label="Composite Score ε - 0.3α")
    for idx, design in enumerate(designs[:5]):
        ax.annotate(str(idx + 1), (design["alpha"], design["epsilon"]))
    fig.tight_layout()
    plt.savefig(path, dpi=300)
    plt.close(fig)


def main():
    designs = optimize()
    outdir = os.path.join(os.path.dirname(__file__), "outputs")
    os.makedirs(outdir, exist_ok=True)
    summary_path = os.path.join(outdir, "question3_multilayer_summary.csv")
    write_summary(designs, summary_path)
    plot_path = os.path.join(outdir, "question3_pareto.png")
    plot_pareto(designs, plot_path)

    print(f"Optimal designs written to: {summary_path}")
    print(f"Performance scatter plot: {plot_path}")
    top = designs[0]
    layer_desc = "; ".join(
        f"{layer['material']}@{layer['thickness']:.2f}um" for layer in top["layers"]
    )
    print(
        "Best design: score={:.3f}, ε={:.3f}, α={:.3f}, layers={}".format(
            top["score"], top["epsilon"], top["alpha"], layer_desc
        )
    )


if __name__ == "__main__":
    main()