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)
        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."""
        df.loc[:, ["lat", "lon", "alt"]] = preprocessors.norm_coords(
            df["lat"].values, df["lon"].values, df["alt"].values
        )
        df.loc[:, "ias"] = preprocessors.norm_ias(df["ias"].values)
        df.loc[:, "timestamp"] = preprocessors.norm_time(df["timestamp"].values)
        for col in ["lat", "lon", "alt", "ias"]:
            df.loc[:, col] = preprocessors.fillna_with_mean(df[col])

        mapped_df = preprocessors.map_categories(df, preprocessors.category_mappings)
        X_a_1 = F.one_hot(torch.tensor(mapped_df["R_1_IDX"].values)).float()
        X_a_2 = F.one_hot(torch.tensor(mapped_df["R_2_IDX"].values)).float()
        X_t = F.one_hot(torch.tensor(mapped_df["T_IDX"].values)).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"]]

        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"]

        return df


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=("r","t"),
        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)
            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[: self.n_input][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=("r","t"),
        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).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=("r","t"),
    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
    

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]