import torch
from torch import nn
import pathlib
import json
import numpy as np
from torch.utils.data import Dataset


MIN_LAT, MAX_LAT = 57.0, 72.0
MIN_LON, MAX_LON = 3.0, 32.0
MIN_ALT, MAX_ALT = -500.0, 50000.0  # in feet
MIN_IAS, MAX_IAS = 0.0, 800.0  # in knots
TIME_NORM_FACTOR = 1/3600.0  # normalize time to hours

def min_max_normalize(value, min_val, max_val):
    return (value - min_val) / (max_val - min_val)

def min_max_denormalize(norm_value, min_val, max_val):
    return norm_value * (max_val - min_val) + min_val

def norm_coords(lat, lon, alt):
    lat_norm = min_max_normalize(lat, MIN_LAT, MAX_LAT)
    lon_norm = min_max_normalize(lon, MIN_LON, MAX_LON)
    alt_norm = min_max_normalize(alt, MIN_ALT, MAX_ALT)
    return lat_norm, lon_norm, alt_norm

def denorm_coords(lat_norm, lon_norm, alt_norm):
    lat = min_max_denormalize(lat_norm, MIN_LAT, MAX_LAT)
    lon = min_max_denormalize(lon_norm, MIN_LON, MAX_LON)
    alt = min_max_denormalize(alt_norm, MIN_ALT, MAX_ALT)
    return lat, lon, alt

def norm_ias(ias):
    return min_max_normalize(ias, MIN_IAS, MAX_IAS)

def denorm_ias(ias_norm):
    return min_max_denormalize(ias_norm, MIN_IAS, MAX_IAS)

def norm_time(timestamp_series):
    invalid_mask = (timestamp_series < 1672527600) | np.isnan(timestamp_series)  # Jan 1, 2023
    if sum(invalid_mask) == len(timestamp_series):
        raise ValueError("All timestamps are invalid or NaN.")
    timestamp_series[invalid_mask] = timestamp_series[~invalid_mask].mean()
    min_time = timestamp_series.min()
    return (timestamp_series - min_time) * TIME_NORM_FACTOR

def fillna_with_mean(series, allow_nan_mean=True):
    mean_value = series.mean()
    if np.isnan(mean_value):
        if not allow_nan_mean:
            raise ValueError("Mean is NaN and allow_nan_mean is False.")
        mean_value = 0.0
    return series.fillna(mean_value)

class EncodingDataset(Dataset):
    def __init__(self, X_a_1, X_a_2, X_b):
        super().__init__()
        self.X_a_1 = X_a_1
        self.X_a_2 = X_a_2
        self.X_b = X_b
    def __len__(self):
        return self.X_a_1.shape[0]
    def __getitem__(self, idx):
        return self.X_a_1[idx], self.X_a_2[idx], self.X_b[idx]

class AutoEncoder(nn.Module):
    def __init__(self, input_a_1_size, input_a_2_size, input_b_size, latent_size):
        super(AutoEncoder, self).__init__()
        self.input_size = input_a_1_size + input_a_2_size + input_b_size
        self.input_a_1_size = input_a_1_size
        self.input_a_2_size = input_a_2_size
        self.input_b_size = input_b_size

        self.encoder = nn.Sequential(
            nn.Linear(self.input_size, 128),
            nn.ReLU(),
            nn.Linear(128, latent_size),
        )
        self.decoder_A_1 = nn.Sequential(
            nn.Linear(latent_size, 128),
            nn.ReLU(),
            nn.Linear(128, self.input_a_1_size),
            nn.Softmax(dim=1),
        )
        self.decoder_A_2 = nn.Sequential(
            nn.Linear(latent_size, 128),
            nn.ReLU(),
            nn.Linear(128, self.input_a_2_size),
            nn.Softmax(dim=1),
        )
        self.decoder_B = nn.Sequential(
            nn.Linear(latent_size, 128),
            nn.ReLU(),
            nn.Linear(128, self.input_b_size),
            nn.Softmax(dim=1),
        )
    def forward(self, X_a_1, X_a_2, X_b):
        X = torch.cat((X_a_1, X_a_2, X_b), dim=1)
        latent = self.encoder(X)
        reconstructed_A_1 = self.decoder_A_1(latent)
        reconstructed_A_2 = self.decoder_A_2(latent)
        reconstructed_B = self.decoder_B(latent)
        return reconstructed_A_1, reconstructed_A_2, reconstructed_B

category_max_limits = {
    "r_1": 100,
    "r_2": 50,
    "t": 100,
}
autoencoder = AutoEncoder(input_a_1_size=category_max_limits["r_1"], input_a_2_size=category_max_limits["r_2"], input_b_size=category_max_limits["t"], latent_size=4)
auto_encoder_path = pathlib.Path(__file__).parent / "autoencoder.pth"
if auto_encoder_path.exists():
    autoencoder.load_state_dict(torch.load(auto_encoder_path))
    autoencoder.eval()

def train_autoencoder(X_a_1, X_a_2, X_b, num_epochs=100, learning_rate=1e-3):
    criterion = nn.CrossEntropyLoss()
    optimizer = torch.optim.Adam(autoencoder.parameters(), lr=learning_rate)
    ds = EncodingDataset(X_a_1, X_a_2, X_b)
    dataloader = torch.utils.data.DataLoader(ds, batch_size=512, shuffle=True)

    # -----------------------------
    # Auto-termination parameters
    # -----------------------------
    patience = 10
    best_loss = float('inf')
    epochs_no_improve = 0

    # -----------------------------
    # Training loop
    # -----------------------------
    for epoch in range(num_epochs):
        autoencoder.train()
        epoch_loss = 0.0

        for Xa1, Xa2, Xb in dataloader:
            optimizer.zero_grad()
            rec_A_1, rec_A_2, rec_B = autoencoder(Xa1, Xa2, Xb)

            # Targets from the one-hot vectors
            yA1 = torch.argmax(Xa1, dim=1)
            yA2 = torch.argmax(Xa2, dim=1)
            yB  = torch.argmax(Xb, dim=1)

            loss = (
                criterion(rec_A_1, yA1) +
                criterion(rec_A_2, yA2) +
                criterion(rec_B,  yB)
            )
            loss.backward()
            optimizer.step()

            epoch_loss += loss.item() * Xa1.size(0)

        epoch_loss /= len(ds)
        print(f"Epoch [{epoch+1}/{num_epochs}] - Loss: {epoch_loss:.4f}")

        # --- Auto-termination ---
        if epoch_loss < best_loss:
            best_loss = epoch_loss
            epochs_no_improve = 0
        else:
            epochs_no_improve += 1
            if epochs_no_improve >= patience:
                print("Training stopped early — loss plateaued.")
                break

    torch.save(autoencoder.state_dict(), auto_encoder_path)

def encode_features(X_a_1, X_a_2, X_b):
    autoencoder.eval()
    with torch.no_grad():
        X = torch.cat((X_a_1, X_a_2, X_b), dim=1)
        latent = autoencoder.encoder(X)
    return latent

mappings_path = pathlib.Path(__file__).parent / "category_mappings.json"


def create_category_mappings(df):
    df = df.loc[:, ["r", "t"]].dropna().reset_index(drop=True)
    df.loc[:, "r_1"] = df["r"].str.split("-").str[0]
    df.loc[:, "r_2"] = df["r"].str.split("-").str[1].fillna("Empty")

    for col, max_unique in category_max_limits.items():
        top_values = df[col].value_counts().head(max_unique-1).index.tolist()
        df.loc[~df[col].isin(top_values), col] = "Other"

    for col in category_max_limits.keys():
        df.loc[:, col.upper()] = df[col].astype("category")
    
    category_mappings = {}
    for col in category_max_limits.keys():
        category_mappings[col] = {cat: idx for idx, cat in enumerate(df[col.upper()].cat.categories)}
    with open(mappings_path, "w") as f:
        json.dump(category_mappings, f)
    return category_mappings

def map_categories(df, category_mappings):
    if category_mappings is None:
        raise ValueError("Category mappings not provided.")
    df = df.loc[:, ["r", "t"]].copy().dropna().reset_index(drop=True)
    df.loc[:, "r_1"] = df["r"].str.split("-").str[0]
    df.loc[:, "r_2"] = df["r"].str.split("-").str[1].fillna("Empty")

    for col in category_mappings.keys():
        mapping = category_mappings[col]
        df.loc[~df[col].isin(mapping.keys()), col] = "Other"
        df.loc[:, f"{col.upper()}_IDX"] = df[col].map(mapping).astype("int64")
    return df


if mappings_path.exists():
    with open(mappings_path, "r") as f:
        category_mappings = json.load(f)
else:
    category_mappings = None
    print("Category mappings file not found. Please create mappings using 'create_category_mappings' function.")