import torch
from torch.autograd import Function
# import modules
from dataclasses import dataclass
from tqdm.auto import tqdm  
import numpy as np
from getdata import GetDataset
# deep learning modules
import scipy.sparse as sp
from scipy.sparse.linalg import spilu
from scipy.io import loadmat
from torch.autograd import Variable
import torch.nn as nn
import torch.nn.functional as F
import torch.optim as optim
from torch.utils.data import Dataset, DataLoader
import datetime
import pandas as pd
# Plot modules
import matplotlib.pyplot as plt
from scipy.interpolate import griddata  
from pathlib import Path
from scipy.io import savemat

device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

class ILUApply(Function):
    @staticmethod
    def forward(ctx, r_torch, ilu):
        """
        r_torch: torch complex tensor, shape (Mi,)
        ilu: fixed SciPy spilu object
        """
        ctx.ilu = ilu
        r_np = r_torch.detach().cpu().numpy()
        z_np = ilu.solve(r_np)  # z = M^{-1} r
        z = torch.from_numpy(z_np).to(r_torch.device).to(r_torch.dtype)
        return z

    @staticmethod
    def backward(ctx, grad_out):
        """
        grad_out: dL/dz
        complex: grad_r = M^{-H} grad_out
        """
        ilu = ctx.ilu
        g_np = grad_out.detach().cpu().numpy()
        gr_np = ilu.solve(g_np, trans='H')  # conjugate-transpose solve
        grad_r = torch.from_numpy(gr_np).to(grad_out.device).to(grad_out.dtype)
        return grad_r, None

@dataclass
class PINNConfig:
    # 训练（仅 PDE 残差时建议 batch_size 较小如 16~32，避免 branch 输出塌缩）
    epochs: int = 1_000
    batch_size: int = 32
    learning_rate: float = 1e-2
    step_size: int = 500  # StepLR: 每 step_size 轮衰减一次
    gamma: float = 0.95  # StepLR 衰减系数
    max_grad_norm: float = 1.0  # 梯度裁剪，稳定训练
    print_every: int = 1  # 每多少轮打印一次
    save_every: int = 1000  # 每多少轮保存一次 checkpoint

    # 数据（n 为样本数目，如 deepOnet_data_A1_100 表示 100 个样本）
    matpath: str = "deepOnet_data_A1_1558_8_2"

    # 保存
    save_dir: str = "./model_save"
    results_dir: str = "./results"
    load_file_name: str = "model_A1_size_1558_2"
    save_file_name: str = "model_A1_size_1558_2"

    # 设备/精度
    device: str = "cuda" if torch.cuda.is_available() else "cpu"
    dtype: torch.dtype = torch.float64  # 和 .mat 的 float64 对齐（double）
    num_workers: int = 10
    pin_memory: bool = True

class Modified_MLP_Block(nn.Module):
    def __init__(self, input_dim, hidden_channel, output_dim, hidden_size=6):
        super(Modified_MLP_Block, self).__init__()
        self.activation = nn.Tanh()
        self.encodeU = nn.Linear(input_dim, hidden_channel)
        self.encodeV = nn.Linear(input_dim, hidden_channel)
        self.In = nn.Linear(input_dim, hidden_channel)

        self.hidden_layers = nn.ModuleList([
            nn.Linear(hidden_channel, hidden_channel) for _ in range(hidden_size)
        ])
        self.out = nn.Linear(hidden_channel, output_dim)
        self._init_weights()

    def _init_weights(self):
        torch.manual_seed(123)
        gain = nn.init.calculate_gain('tanh')
        for m in self.modules():
            if isinstance(m, nn.Linear):
                nn.init.xavier_uniform_(m.weight, gain=gain)
                nn.init.zeros_(m.bias)

    def forward(self, x):
        U = self.activation(self.encodeU(x))
        V = self.activation(self.encodeV(x))
        Hidden = self.activation(self.In(x))

        for layer in self.hidden_layers:
            Z = self.activation(layer(Hidden))
            Hidden = (1 - Z) * U + Z * V

        x = self.out(Hidden)
        return x
    


def _branch_norm2d(channels):
    """Branch 内使用 InstanceNorm2d：按样本、按通道独立归一化，不跨样本混合，避免输出塌缩。"""
    return nn.InstanceNorm2d(channels)


def add_spatial_coord_channels(epsilon_data):
    """
    在 ε 图像上拼接空间坐标通道，使 CNN 能区分「同形不同位」的介质。
    epsilon_data: (B, 1, H, W) -> 返回 (B, 3, H, W)，通道为 [ε, x_norm, y_norm]，归一化到 [0,1]。
    """
    B, _, H, W = epsilon_data.shape
    device, dtype = epsilon_data.device, epsilon_data.dtype
    j = torch.linspace(0, 1, W, device=device, dtype=dtype).view(1, 1, 1, W).expand(B, 1, H, W)
    i = torch.linspace(0, 1, H, device=device, dtype=dtype).view(1, 1, H, 1).expand(B, 1, H, W)
    return torch.cat([epsilon_data, j, i], dim=1)


class CNN_Branch_Residual(nn.Module):
    """带残差连接的 CNN 分支（Branch 用 InstanceNorm2d；输入含 ε + 空间坐标通道以区分同形不同位）。"""

    def __init__(self, in_channels=3, num_classes=128):
        super(CNN_Branch_Residual, self).__init__()

        # 初始卷积层
        self.initial = nn.Sequential(
            nn.Conv2d(in_channels, 32, kernel_size=3, padding=1, bias=False),
            _branch_norm2d(32),
            nn.ReLU(inplace=True),
            nn.Conv2d(32, 32, kernel_size=3, padding=1, bias=False),
            _branch_norm2d(32),
            nn.ReLU(inplace=True)
        )

        # 残差块
        self.res_block1 = ResidualBlock(32, 64, stride=2, norm_layer=_branch_norm2d)
        self.res_block2 = ResidualBlock(64, 128, stride=2, norm_layer=_branch_norm2d)

        # 全局平均池化
        self.global_avg_pool = nn.AdaptiveAvgPool2d((1, 1))

        # 全连接层
        self.fc = nn.Linear(128, num_classes)

    def forward(self, x):
        x = self.initial(x)
        x = self.res_block1(x)
        x = self.res_block2(x)
        x = self.global_avg_pool(x)
        x = x.view(x.size(0), -1)
        x = self.fc(x)
        return x


class ResidualBlock(nn.Module):
    """残差块（支持 norm_layer，Branch 中传入 GroupNorm）。"""

    def __init__(self, in_channels, out_channels, stride=1, norm_layer=None):
        super(ResidualBlock, self).__init__()
        if norm_layer is None:
            norm_layer = nn.BatchNorm2d

        self.conv1 = nn.Conv2d(in_channels, out_channels, kernel_size=3,
                               stride=stride, padding=1, bias=False)
        self.bn1 = norm_layer(out_channels)

        self.conv2 = nn.Conv2d(out_channels, out_channels, kernel_size=3,
                               stride=1, padding=1, bias=False)
        self.bn2 = norm_layer(out_channels)

        self.relu = nn.ReLU(inplace=True)

        # 下采样连接
        self.downsample = None
        if stride != 1 or in_channels != out_channels:
            self.downsample = nn.Sequential(
                nn.Conv2d(in_channels, out_channels, kernel_size=1,
                         stride=stride, bias=False),
                norm_layer(out_channels)
            )

    def forward(self, x):
        identity = x

        out = self.conv1(x)
        out = self.bn1(out)
        out = self.relu(out)

        out = self.conv2(out)
        out = self.bn2(out)

        if self.downsample is not None:
            identity = self.downsample(x)

        out += identity
        out = self.relu(out)

        return out

class DeepONet(nn.Module):
    def __init__(self, branch_input_dim, trunk_input_dim, hidden_channel, output_dim):
        super(DeepONet, self).__init__()
        self.output_dim = output_dim
        # branch 输入为 ε + 2 个空间坐标通道，共 3 通道，便于区分同形不同位的介质
        self.branch_net = CNN_Branch_Residual(in_channels=3, num_classes=output_dim)
        self.trunk_net = Modified_MLP_Block(trunk_input_dim, hidden_channel, output_dim)

    def forward(self, branch_input, trunk_input):
        branch_input = add_spatial_coord_channels(branch_input)
        branch_out = self.branch_net(branch_input)
        trunk_out = self.trunk_net(trunk_input)

        B1, B2 = branch_out[:, :self.output_dim//2], branch_out[:, self.output_dim//2:]
        T1, T2 = trunk_out[:, :, :self.output_dim//2], trunk_out[:, :, self.output_dim//2:]
        #print("B1 shape:", B1.shape, "B2 shape:", B2.shape)
        #print("T1 shape:", T1.shape, "T2 shape:", T2.shape)
        s_re = torch.einsum('bi,bni->bn', B1, T1) #实部
        s_im = torch.einsum('bi,bni->bn', B2, T2)
        return s_re, s_im


class PINN_maxwell():
    def __init__(self, model, config: PINNConfig):
        self.cfg = config

        self.device = torch.device(self.cfg.device)
        self.model = model.to(self.device, dtype=self.cfg.dtype)
        self.batch_size = self.cfg.batch_size
        self.learning_rate = self.cfg.learning_rate
        self.matpath = self.cfg.matpath
        self.loss_fn = nn.MSELoss()
        self.optimizer = optim.Adam(self.model.parameters(), lr=self.learning_rate)
        self.scheduler = optim.lr_scheduler.StepLR(self.optimizer, step_size=self.cfg.step_size, gamma=self.cfg.gamma)
        self.losses = []
        self.lamda = []
        self.save_file_name = self.cfg.save_file_name
        self.load_file_name = self.cfg.load_file_name
        self.save_dir = Path(self.cfg.save_dir)
        self.train_set, self.test_set = self.load_dataset()
        self.train_loader = DataLoader(self.train_set, self.cfg.batch_size, shuffle=True)
        self.test_loader = DataLoader(self.test_set, batch_size=len(self.test_set), shuffle=False)
        self.ilu_cache = {}
        self._build_ilu_cache()
        

    def load_model(self):
        self.model.load_state_dict(torch.load(self.save_dir / f'{self.load_file_name}.pth', map_location=self.device, weights_only=True))

    def E_function(self, epsilon_data, coord_data):
        epsilon_data = epsilon_data.to(self.device)
        coord_data = coord_data.to(self.device)
        return self.model(epsilon_data, coord_data)
    
    def get_data_loss(self, epsilon_data, coord_data, E_true):
        E_re_pred, E_im_pred = self.E_function(epsilon_data, coord_data)
        E_re_true = E_true[:,:, 0]
        E_im_true = E_true[:,:, 1]
        data_loss = self.loss_fn(E_re_pred, E_re_true) + self.loss_fn(E_im_pred, E_im_true)
        return data_loss

    def get_fem_loss(self, indices, epsilon_data, coord_data, E_true, Ai, Aj, Av, b, coord_len):
        """indices: (B,) 每个样本的全局索引，用于从 ilu_cache 取对应的 ILU。"""
        Ere_pred, Eim_pred = self.E_function(epsilon_data, coord_data)
        E = torch.complex(Ere_pred, Eim_pred)  # shape: (B, Mmax)

        B, Mmax = E.shape
        Mi = coord_len.squeeze(-1).long().to(self.device)  # shape: (B,)

        arangeM = torch.arange(Mmax, device=self.device)  # (Mmax,)
        mask_x = arangeM[None, :] < Mi[:, None]  # (B, Mmax)

        x_flat = E[mask_x]  # (sum Mi,)
        b_flat = b.to(self.device)[mask_x].to(x_flat.dtype)  # (sum Mi,)

        sumMi = int(Mi.sum().item())
        offsets = torch.cumsum(
            torch.cat([torch.zeros(1, device=self.device, dtype=torch.long), Mi[:-1]]),
            dim=0
        )  # (B,)

        Ai = Ai.to(self.device).long()
        Aj = Aj.to(self.device).long()
        Av = Av.to(self.device).to(x_flat.dtype)
        mask_nnz = (Ai > 0) & (Aj > 0)  # (B, Kmax)，padding 位置为 0

        rows = (Ai - 1 + offsets.unsqueeze(1)).masked_select(mask_nnz)
        cols = (Aj - 1 + offsets.unsqueeze(1)).masked_select(mask_nnz)
        vals = Av.masked_select(mask_nnz)

        y = torch.zeros(sumMi, dtype=x_flat.dtype, device=self.device)
        y.scatter_add_(0, rows, vals * x_flat.index_select(0, cols))
        r = y - b_flat

        # 按样本用缓存的 ILU 计算 z，每个样本一个 ILU
        z_parts = []
        for i in range(B):
            start = int(offsets[i].item())
            m = int(Mi[i].item())
            r_i = r[start : start + m]
            ilu = self.ilu_cache[int(indices[i].item())]
            z_i = ILUApply.apply(r_i, ilu)
            z_parts.append(z_i)
        z = torch.cat(z_parts, dim=0)

        loss = (z.abs() ** 2).mean()
        return loss
    
    
    @torch.no_grad()
    def test_E_loss(self):
        self.model.eval()
        total_loss = 0.0
        for indices, epsilon_data, coord_data, E_true, Ai, Aj, Av, b, coord_len in self.test_loader:
            indices = indices.to(self.device)
            epsilon_data = epsilon_data.to(self.device)
            coord_data = coord_data.to(self.device)
            loss = self.get_fem_loss(indices, epsilon_data, coord_data, E_true, Ai, Aj, Av, b, coord_len)
            total_loss += loss.item()
        self.model.train()
        return total_loss / len(self.test_loader)
    
    def train(self, epochs, print_every=100, save_every=10000):
        self.losses.append(['epoch', 'fem_loss', 'test_loss'])
        start_time = datetime.datetime.now()
        for epoch in tqdm(range(epochs), desc='Training'):
            self.model.train()
            total_loss = 0.0
            fem_loss = 0.0

            for indices, epsilon_data, coord_data, E_true, Ai, Aj, Av, b, coord_len in self.train_loader:
                indices = indices.to(self.device)
                epsilon_data = epsilon_data.to(self.device)
                coord_data = coord_data.to(self.device)
                self.optimizer.zero_grad()

                fem_loss = self.get_fem_loss(indices, epsilon_data, coord_data, E_true, Ai, Aj, Av, b, coord_len)

                loss = fem_loss

                loss.backward()
                if self.cfg.max_grad_norm > 0:
                    torch.nn.utils.clip_grad_norm_(self.model.parameters(), self.cfg.max_grad_norm)
                self.optimizer.step()

                total_loss += loss.item()

            avg_total_loss = total_loss / len(self.train_loader)
            avg_test_loss = self.test_E_loss()
            self.losses.append([epoch, avg_total_loss, avg_test_loss])
            self.scheduler.step()

            if epoch % print_every == 0:
                print(f'Epoch {epoch}, Total Loss: {avg_total_loss}, test Loss {avg_test_loss}')
            if (epoch + 1) % save_every == 0:
                ckpt_path = self.save_dir / f'{self.save_file_name}_epoch{epoch + 1}.pth'
                torch.save(self.model.state_dict(), ckpt_path)
        torch.save(self.model.state_dict(), self.save_file_name)
        print("Current learning rate:", self.optimizer.param_groups[0]['lr'])
        print("Training Time:", (datetime.datetime.now() - start_time).total_seconds(), "s")

    def plot_loss(self):
        data = np.array(self.losses[1:])
        epochs = data[:, 0]
        train_loss = data[:, 1]
        test_loss = data[:, 2]
        plt.figure(figsize=(10, 6))
        plt.title('Training/Test Loss')
        plt.semilogy(epochs, train_loss, label='train_loss')
        plt.semilogy(epochs, test_loss,  label='test_loss')
        plt.xlabel('Epochs')
        plt.ylabel('Loss')
        plt.legend()
        plt.grid()
        save_path = 'results/loss_plot2_1.png'
        plt.savefig(save_path)
        plt.show()
    
    def load_dataset(self):
        """从 MATLAB 划分好的 .mat 文件加载数据，90% 训练 10% 测试由 mat 内已分好。"""
        data_set = loadmat(self.matpath)

        Epsilon_train = data_set['Eplison_train']
        X_train = data_set['X_train']
        Ez_train = data_set['Ez_train']

        Epsilon_test = data_set['Eplison_test']
        X_test = data_set['X_test']
        Ez_test = data_set['Ez_test']

        coord_len_train = data_set['coord_len_train']
        coord_len_test = data_set['coord_len_test']

        Ai_train, Aj_train = data_set['Ai_train'], data_set['Aj_train']
        Ai_test, Aj_test = data_set['Ai_test'], data_set['Aj_test']
        Av_train, Av_test = data_set['Av_train'], data_set['Av_test']
        b_train, b_test = data_set['b_train'], data_set['b_test']

        n_train, n_test = len(Epsilon_train), len(Epsilon_test)
        print(f"Train samples: {n_train}, Test samples: {n_test}")
        print(f"Train shapes: ε {Epsilon_train.shape}, X {X_train.shape}, Ez {Ez_train.shape}")

        Train_dataset = GetDataset(
            Epsilon_train, X_train, Ez_train,
            Ai_train, Aj_train, Av_train, b_train, coord_len_train,
            index_offset=0
        )
        Test_dataset = GetDataset(
            Epsilon_test, X_test, Ez_test,
            Ai_test, Aj_test, Av_test, b_test, coord_len_test,
            index_offset=n_train
        )
        return Train_dataset, Test_dataset

    def _build_ilu_cache(self):
        """训练开始前为每个样本的稀疏矩阵 A 预计算 ILU 并缓存，训练过程中 A 不变，只算一次。"""
        for idx in tqdm(range(len(self.train_set)), desc="Building ILU cache (train)"):
            ai = self.train_set.Ai[idx].numpy()
            aj = self.train_set.Aj[idx].numpy()
            av = self.train_set.Av[idx].numpy()
            Mi = int(self.train_set.coord_len[idx].item())
            mask = (ai > 0) & (aj > 0)
            rows = (ai[mask] - 1).astype(np.int64)
            cols = (aj[mask] - 1).astype(np.int64)
            vals = av[mask]
            A = sp.coo_matrix((vals, (rows, cols)), shape=(Mi, Mi)).tocsc()
            self.ilu_cache[idx] = spilu(A)
        for idx in tqdm(range(len(self.test_set)), desc="Building ILU cache (test)"):
            ai = self.test_set.Ai[idx].numpy()
            aj = self.test_set.Aj[idx].numpy()
            av = self.test_set.Av[idx].numpy()
            Mi = int(self.test_set.coord_len[idx].item())
            mask = (ai > 0) & (aj > 0)
            rows = (ai[mask] - 1).astype(np.int64)
            cols = (aj[mask] - 1).astype(np.int64)
            vals = av[mask]
            A = sp.coo_matrix((vals, (rows, cols)), shape=(Mi, Mi)).tocsc()
            self.ilu_cache[len(self.train_set) + idx] = spilu(A)
        print(f"ILU cache built: {len(self.ilu_cache)} samples.")

    def saveE_pred(self):
        for indices, epsilon_data, coord_data, *_ in self.train_loader:
            epsilon_data = epsilon_data.to(self.device)
            coord_data = coord_data.to(self.device)
            E_real, E_imag = self.E_function(epsilon_data, coord_data)
            E_pred = torch.complex(E_real, E_imag)
        E_pred_np = E_pred.detach().cpu().numpy()
        savemat('E_train_pred_size_1588.mat', {"E_pred": E_pred_np})
        for indices, epsilon_data, coord_data, *_ in self.test_loader:
            epsilon_data = epsilon_data.to(self.device)
            coord_data = coord_data.to(self.device)
            E_real, E_imag = self.E_function(epsilon_data, coord_data)
            E_pred = torch.complex(E_real, E_imag)
        E_pred_np = E_pred.detach().cpu().numpy()
        savemat('E_test_pred_size_1588.mat', {"E_pred": E_pred_np})


if __name__ == "__main__":
    cfg = PINNConfig(
    )
    model = DeepONet(branch_input_dim=1, trunk_input_dim=2, hidden_channel=128, output_dim=64)
    model = model.double()
    pinn = PINN_maxwell(model, cfg)
    # pinn.load_model()
    #pinn.test_fem_loss()
    pinn.train(epochs=cfg.epochs, print_every=cfg.print_every, save_every=cfg.save_every)
    pinn.plot_loss()
    # pinn.saveE_pred()