import pandas as pd
import torch
import torch.nn.functional as F
from torch.utils.data import Dataset, DataLoader
import numpy as np
import pathlib
from typing import Sequence
import random
from aiRNN import preprocessors

class BaseDataset(Dataset):
    def __init__(
        self,
        filepaths,
        n_input,
        n_output,
        n_windows_per_file,
        step=1,
        feature_columns=("lat", "lon", "alt", "ias"),
        context_columns=(*[f"type_encoding_{i}" for i in range(4)], "last_lat", "last_lon", "last_alt", "last_ias"),
        time_columns=("timestamp",),
        target_columns=("lat", "lon", "alt"),
        device="cpu",
    ):
        super().__init__()
        self.filepaths = [pathlib.Path(fp) for fp in filepaths]
        self.n_input = n_input
        self.n_output = n_output
        self.window_size = n_input + n_output
        self.device = device
        self.n_windows_per_file = n_windows_per_file
        self.step = step

        # Which columns to use
        sample = pd.read_csv(self.filepaths[0], nrows=5)
        sample = self._modify_df(sample)
        self.feature_cols = list(feature_columns)
        self.context_cols = list(context_columns)
        self.time_cols = list(time_columns)
        self.target_cols = list(target_columns)

        # Total samples = windows_per_file * number_of_files
        self.total_windows = n_windows_per_file * len(self.filepaths)
        for col_list in [self.feature_cols, self.context_cols, self.time_cols, self.target_cols]:
            assert all(
                col in sample.columns for col in col_list
            ), "Some specified columns not found in data."

    def __len__(self):
        return self.total_windows
    
    def _get_start_end_indices(self, idx, n_rows):
        # Max valid start index so window stays inside the file
        max_start = n_rows - self.window_size * self.step
        if max_start < 0:
            raise ValueError("File too small for one window.")

        # Evenly spaced index:
        # raw_start = (n_rows / n_windows_per_file) * local_idx
        # mapped into integer space
        raw_start = (n_rows / self.n_windows_per_file) * idx
        start_idx = int(raw_start // self.step) * self.step

        # Clamp to safe zone
        if start_idx > max_start:
            start_idx = max_start

        end_idx = start_idx + self.window_size * self.step
        return start_idx, end_idx

    def _modify_df(self, df):
        """Hook for subclasses to modify dataframe before slicing windows."""
        lat_norm, lon_norm, alt_norm = preprocessors.norm_coords(
            df["lat"].values, df["lon"].values, df["alt"].values
        )
        df.loc[:, "lat"] = lat_norm
        df.loc[:, "lon"] = lon_norm
        df.loc[:, "alt"] = alt_norm
        df.loc[:, "ias"] = preprocessors.norm_ias(df["ias"].values)
        df.loc[:, "dt"] = df["timestamp"].diff()
        df.loc[:, "dt"] = preprocessors.fillna_with_mean(df["dt"], allow_nan_mean=False)
        df.loc[:, "timestamp"] = preprocessors.norm_time(df["timestamp"].values)
        for col in ["lat", "lon", "alt", "ias"]:
            if df[col].isnull().all() and col != "ias":
                raise ValueError(f"All values in column {col} are NaN.")
            df.loc[:, col] = preprocessors.fillna_with_mean(df[col], allow_nan_mean=(col == "ias"))

        mapped_df = preprocessors.map_categories(df, preprocessors.category_mappings)
        max_values = list(preprocessors.category_max_limits.values())
        X_a_1 = F.one_hot(torch.tensor(mapped_df["R_1_IDX"].values), max_values[0]).float()
        X_a_2 = F.one_hot(torch.tensor(mapped_df["R_2_IDX"].values), max_values[1]).float()
        X_t = F.one_hot(torch.tensor(mapped_df["T_IDX"].values), max_values[2]).float()
        df.loc[:, [f"type_encoding_{i}" for i in range(4)]] = preprocessors.encode_features(
            X_a_1, X_a_2, X_t
        )
        last_row = df.iloc[-1][["lat", "lon", "alt", "ias", "timestamp"]]

        df.loc[:, "last_lat"] = last_row["lat"]
        df.loc[:, "last_lon"] = last_row["lon"]
        df.loc[:, "last_alt"] = last_row["alt"]
        df.loc[:, "last_ias"] = last_row["ias"]
        df.loc[:, "last_timestamp"] = last_row["timestamp"]

        return df
    
    def save_entire_dataset(self, filepath):
        """Utility to save the entire dataset to a single pth file."""
        X_feat_list = []
        X_time_list = []
        Y_out_list = []
        X_context_list = []
        for i in range(len(self)):
            X_f, X_t, Y, X_c = self[i]
            X_feat_list.append(X_f.cpu())
            X_time_list.append(X_t.cpu())
            Y_out_list.append(Y.cpu())
            if X_c is not None:
                X_context_list.append(X_c.cpu())
        data_dict = {
            "X_feat": torch.stack(X_feat_list),
            "X_time": torch.stack(X_time_list),
            "Y_out": torch.stack(Y_out_list),
        }
        if len(X_context_list) > 0:
            data_dict["X_context"] = torch.stack(X_context_list)
        torch.save(data_dict, filepath)

class SaveDataset(Dataset):
    def __init__(self, data_dict, device="cpu"):
        super().__init__()
        self.X_feat = data_dict["X_feat"].to(device)
        self.X_time = data_dict["X_time"].to(device)
        self.Y_out = data_dict["Y_out"].to(device)
        self.device = device
        if "X_context" in data_dict:
            self.X_context = data_dict["X_context"].to(device)
        else:
            self.X_context = None

    def __len__(self):
        return self.X_feat.shape[0]

    def __getitem__(self, idx):
        X_f = self.X_feat[idx]
        X_t = self.X_time[idx]
        Y = self.Y_out[idx]
        if self.X_context is not None:
            X_c = self.X_context[idx]
        else:
            X_c = None
        return X_f, X_t, Y, X_c

class EvenlySpacedDataset(BaseDataset):
    def __init__(
        self,
        filepaths,
        n_input,
        n_output,
        n_windows_per_file,
        step=1,
        feature_columns=("lat", "lon", "alt", "ias"),
        context_columns=(*[f"type_encoding_{i}" for i in range(4)], "last_lat", "last_lon", "last_alt", "last_ias"),
        time_columns=("timestamp",),
        target_columns=("lat", "lon", "alt"),
        device="cpu",
    ):
        super().__init__(
            filepaths=filepaths,
            n_input=n_input,
            n_output=n_output,
            feature_columns=feature_columns,
            context_columns=context_columns,
            time_columns=time_columns,
            target_columns=target_columns,
            device=device,
            step=step,
            n_windows_per_file=n_windows_per_file,
        )

        # Preload all data into memory
        self.data = []
        for fp in self.filepaths:
            df = pd.read_csv(fp, dtype={"icao": str, "r": str, "t": str, "timestamp": float, "lat": float, "lon": float, "alt": float, "ias": float})
            n_rows = len(df)
            # Extract windows
            for w in range(n_windows_per_file):
                try:
                    start_idx, end_idx = self._get_start_end_indices(w, n_rows)
                    window = df.iloc[start_idx:end_idx:self.step]
                    window = self._modify_df(window)
                    self.data.append(window)
                except ValueError:
                    self.total_windows -= 1
                    continue

    def __len__(self):
        return self.total_windows

    def __getitem__(self, idx):
        window = self.data[idx]

        X_feat = window.iloc[: self.n_input][self.feature_cols].to_numpy(np.float32)
        if len(self.context_cols) > 0:
            X_context = window.iloc[0][self.context_cols].to_numpy(np.float32)
        else:
            X_context = None
        X_time = window.iloc[: self.window_size][self.time_cols].to_numpy(np.float32)
        Y_out = window.iloc[self.n_input :][self.target_cols].to_numpy(np.float32)
        return (
            torch.tensor(X_feat, device=self.device),
            torch.tensor(X_time, device=self.device),
            torch.tensor(Y_out, device=self.device),
            torch.tensor(X_context, device=self.device) if X_context is not None else None,
        )

class EvenlySpacedStreamingDataset(BaseDataset):
    def __init__(
        self,
        filepaths,
        n_input,
        n_output,
        n_windows_per_file,
        step=1,
        feature_columns=("lat", "lon", "alt", "ias"),
        context_columns=(*[f"type_encoding_{i}" for i in range(4)], "last_lat", "last_lon", "last_alt", "last_ias"),
        time_columns=("timestamp",),
        target_columns=("lat", "lon", "alt"),
        device="cpu",
    ):
        super().__init__(
            filepaths=filepaths,
            n_input=n_input,
            n_output=n_output,
            feature_columns=feature_columns,
            context_columns=context_columns,
            time_columns=time_columns,
            target_columns=target_columns,
            device=device,
            step=step,
            n_windows_per_file=n_windows_per_file,
        )
        
        # Precompute row counts for each file
        self.row_counts = [self._count_rows(fp) for fp in self.filepaths]


    def _count_rows(self, fp):
        with open(fp, "r") as f:
            return sum(1 for _ in f) - 1  # minus header

    def __len__(self):
        return self.total_windows

    def __getitem__(self, idx):
        # Figure out which file we belong to
        file_idx = idx // self.n_windows_per_file
        local_idx = idx % self.n_windows_per_file
        fp = self.filepaths[file_idx]

        n_rows = self.row_counts[file_idx]

        start_idx, end_idx = self._get_start_end_indices(local_idx, n_rows)

        # Read only the required rows
        skip = list(set(range(1, start_idx + 1)))
        df = pd.read_csv(fp, skiprows=skip, nrows=self.window_size * self.step, dtype={"icao": str, "r": str, "t": str, "timestamp": float, "lat": float, "lon": float, "alt": float, "ias": float}).iloc[::self.step]
        df = self._modify_df(df)

        # Slice into input / decoder-input / targets
        X_feat = df.iloc[: self.n_input][self.feature_cols].to_numpy(np.float32)
        if len(self.context_cols) > 0:
            X_context = df.iloc[: self.n_input][self.context_cols].to_numpy(np.float32)
        else:
            X_context = None
        X_time = df.iloc[: self.window_size][self.time_cols].to_numpy(np.float32)
        Y_out = df.iloc[self.n_input :][self.target_cols].to_numpy(np.float32)
        return (
            torch.tensor(X_feat, device=self.device),
            torch.tensor(X_time, device=self.device),
            torch.tensor(Y_out, device=self.device),
            torch.tensor(X_context, device=self.device) if X_context is not None else None,
        )

def get_datasets(
    files: Sequence[str],
    n_input: int,
    n_output: int,
    n_windows_per_file: int,
    step: int = 1,
    feature_columns=("lat", "lon", "alt", "ias"),
    context_columns=(*[f"type_encoding_{i}" for i in range(4)], "last_lat", "last_lon", "last_alt", "last_ias"),
    time_columns=("timestamp",),
    target_columns=("lat", "lon", "alt"),
    device="cpu",
    test_split: float = 0.2,
    val_split: float = 0.1,
    seed: int = 42,
    streaming: bool = False,
) -> tuple[Dataset, Dataset, Dataset]:
    """Utility to create train/val/test datasets from file list."""
    random.seed(seed)
    files = list(files)
    random.shuffle(files)

    n_total = len(files)
    n_test = int(n_total * test_split)
    n_val = int(n_total * val_split)

    test_files = files[:n_test]
    val_files = files[n_test : n_test + n_val]
    train_files = files[n_test + n_val :]

    if not streaming:
        ds = EvenlySpacedDataset
    else:
        ds = EvenlySpacedStreamingDataset

    train_dataset = ds(
        train_files,
        n_input=n_input,
        n_output=n_output,
        n_windows_per_file=n_windows_per_file,
        step=step,
        feature_columns=feature_columns,
        context_columns=context_columns,
        time_columns=time_columns,
        target_columns=target_columns,
        device=device,
    )
    val_dataset = ds(
        val_files,
        n_input=n_input,
        n_output=n_output,
        n_windows_per_file=n_windows_per_file,
        step=step,
        feature_columns=feature_columns,
        context_columns=context_columns,
        time_columns=time_columns,
        target_columns=target_columns,
        device=device,
    )
    test_dataset = ds(
        test_files,
        n_input=n_input,
        n_output=n_output,
        n_windows_per_file=n_windows_per_file,
        step=step,
        feature_columns=feature_columns,
        context_columns=context_columns,
        time_columns=time_columns,
        target_columns=target_columns,
        device=device,
    )

    return train_dataset, val_dataset, test_dataset