import os
import numpy as np
import torch # type: ignore
import torch.nn as nn # type: ignore
import torch.nn.functional as F # type: ignore

os.environ.setdefault('KMP_DUPLICATE_LIB_OK', 'TRUE')

VOLTAGE_SPACE = np.linspace(-1.0, 0.0, 99)
VOLTAGE = VOLTAGE_SPACE[VOLTAGE_SPACE <= -0.40]
N_SIG = len(VOLTAGE)
N_INT_TS = 3 
N_INT_FEAT = 20 

VOLTAGE_INTERVALS = [
    ('int1', -1.00, -0.82),
    ('int2', -0.82, -0.62),
    ('int3', -0.62, -0.40),
]

FEATURE_SUFFIXES = [
    'valley',
    'kurtosis',
    'skewness', 
    'area',
    'valley_position', 
    'peak_width',
    'd1_max', 
    'd1_min', 
    'n_zero_crossings',
    'd2_max', 
    'd2_min',
    'mean', 
    'std', 
    'range', 
    'energy',
    'valley_to_mean',
    'asymmetry',
    'slope_start', 
    'slope_end', 
    'overall_slope',
]

INTERVAL_LABELS = [
    f"{v0:.2f}–{v1:.2f} V"
    for _, v0, v1 in VOLTAGE_INTERVALS
]

FEAT_NAMES = [
    f"int_({v0:.2f},{v1:.2f})_{feature}"
    for _, v0, v1 in VOLTAGE_INTERVALS
    for feature in FEATURE_SUFFIXES
]

INTERVAL_COLORS = ['#1f77b4', '#ff7f0e', '#2ca02c']

class MLPNet(nn.Module):
    def __init__(self, n_features, num_classes):
        super().__init__()
        self.fc1 = nn.Linear(n_features, 64)
        self.bn1 = nn.BatchNorm1d(64)
        self.drop1 = nn.Dropout(0.3)
        self.fc2 = nn.Linear(64, 32)
        self.bn2 = nn.BatchNorm1d(32)
        self.drop2 = nn.Dropout(0.3)
        self.fc3 = nn.Linear(32, 16)
        self.fc_out = nn.Linear(16, num_classes)

    def forward(self, x):
        x = self.drop1(F.relu(self.bn1(self.fc1(x))))
        x = self.drop2(F.relu(self.bn2(self.fc2(x))))
        return self.fc_out(F.relu(self.fc3(x)))


class LSTMWithAttn(nn.Module):
    def __init__(self, n_features, num_classes, hidden=64):
        super().__init__()
        self.lstm    = nn.LSTM(n_features, hidden, batch_first=True, bidirectional=True)
        self.norm    = nn.LayerNorm(hidden * 2)
        self.drop    = nn.Dropout(0.3)
        self.attn    = nn.Linear(hidden * 2, 1)
        self.fc1     = nn.Linear(hidden * 2, 32)
        self.drop_fc = nn.Dropout(0.2)
        self.fc_out  = nn.Linear(32, num_classes)

    def forward(self, x):
            x, _= self.lstm(x)
            x= self.drop(self.norm(x))
            weights = torch.softmax(self.attn(x), dim=1)
            pooled  = (weights * x).sum(dim=1)
            return self.fc_out(self.drop_fc(F.relu(self.fc1(pooled))))


class DualBranchLSTM(nn.Module):
    def __init__(self, n_sig_features, n_int_features, num_classes, hidden_a=64, hidden_b=32):
        super().__init__()
        self.lstm_a = nn.LSTM(n_sig_features, hidden_a, batch_first=True, bidirectional=True)
        self.norm_a = nn.LayerNorm(hidden_a * 2)
        self.drop_a = nn.Dropout(0.3)
        self.attn_a = nn.Linear(hidden_a * 2, 1)

        self.lstm_b = nn.LSTM(n_int_features, hidden_b, batch_first=True, bidirectional=True)
        self.norm_b = nn.LayerNorm(hidden_b * 2)
        self.drop_b = nn.Dropout(0.2)
        self.attn_b = nn.Linear(hidden_b * 2, 1)

        fused_dim = hidden_a * 2 + hidden_b * 2
        self.norm_fuse = nn.LayerNorm(fused_dim)
        self.drop_fuse = nn.Dropout(0.3)
        self.fc1 = nn.Linear(fused_dim, 32)
        self.fc_out = nn.Linear(32, num_classes)

    def forward(self, x_signal, x_intervals):
        a, _ = self.lstm_a(x_signal)
        a = self.drop_a(self.norm_a(a))
        weights = torch.softmax(self.attn_a(a), dim=1)
        a = (weights * a).sum(dim=1)

        b, _ = self.lstm_b(x_intervals)
        b = self.drop_b(self.norm_b(b))
        weights_b = torch.softmax(self.attn_b(b), dim=1)
        b = (weights_b * b).sum(dim=1)

        x = torch.cat([a, b], dim=1)
        x = self.drop_fuse(self.norm_fuse(x))
        return self.fc_out(F.relu(self.fc1(x)))


class MetaLearner(nn.Module):
    def __init__(self, n_base_models, num_classes):
        super().__init__()
        self.fc1 = nn.Linear(n_base_models * num_classes, 32)
        self.drop = nn.Dropout(0.3)
        self.fc2 = nn.Linear(32, num_classes)

    def forward(self, x):
        return self.fc2(self.drop(F.relu(self.fc1(x))))


class FlatWrapper(nn.Module):
    def __init__(self, model, T, F):
        super().__init__()
        self.model = model
        self.T = T
        self.F = F

    def forward(self, x):
        return self.model(x.reshape(-1, self.T, self.F))


class DualInputWrapper(nn.Module):
    def __init__(self, model, n_sig, n_int_ts, n_int_feat):
        super().__init__()
        self.model = model
        self.n_sig = n_sig
        self.n_int_ts = n_int_ts
        self.n_int_feat = n_int_feat

    def forward(self, x):
        sig = x[:, :self.n_sig].unsqueeze(-1)
        intervals = x[:, self.n_sig:].reshape(-1, self.n_int_ts, self.n_int_feat)
        return self.model(sig, intervals)