"""
ATFMTraj Data Loading and Preprocessing for LLM4AirTrack.

Loads ENU-transformed ADS-B trajectories from petchthwr/ATFMTraj.
Creates sliding-window samples: [context_window] -> [prediction_horizon].
Computes kinematic features: directional vectors, polar components, speed proxies.
"""

import os
import numpy as np
import pandas as pd
import torch
from torch.utils.data import Dataset, DataLoader
from huggingface_hub import hf_hub_download
from typing import Tuple, Optional


def download_atfm_dataset(airport="RKSIa", cache_dir="./data/ATFMTraj"):
    """Download ATFMTraj TSV files from HuggingFace Hub."""
    os.makedirs(cache_dir, exist_ok=True)
    airport_dir = os.path.join(cache_dir, airport)
    os.makedirs(airport_dir, exist_ok=True)
    for mode in ["TRAIN", "TEST"]:
        for var in ["X", "Y", "Z"]:
            fname = f"{airport}_{mode}_{var}.tsv"
            fpath = os.path.join(airport_dir, fname)
            if not os.path.exists(fpath):
                print(f"Downloading {airport}/{fname}...")
                hf_hub_download(
                    repo_id="petchthwr/ATFMTraj",
                    filename=f"{airport}/{fname}",
                    repo_type="dataset",
                    local_dir=cache_dir,
                )
    return airport_dir


def load_atfm_raw(airport="RKSIa", mode="TRAIN", cache_dir="./data/ATFMTraj"):
    """Load raw ATFMTraj data. Returns (N, T_max, 3) ENU + (N,) labels."""
    airport_dir = os.path.join(cache_dir, airport)
    data, labels = [], None
    for var in ['X', 'Y', 'Z']:
        df = pd.read_csv(
            os.path.join(airport_dir, f"{airport}_{mode}_{var}.tsv"),
            sep='\t', header=None, na_values='NaN'
        )
        if labels is None:
            labels = df.values[:, 0]
        data.append(df.values[:, 1:])
    return np.stack(data, axis=-1), labels.astype(int)


def compute_kinematic_features(trajectory, dt=1.0):
    """
    Compute 9-dim kinematic features from ENU (x,y,z):
    Position (x,y,z) + Direction (ux,uy,uz) + Polar (r, sinθ, cosθ)
    """
    x, y, z = trajectory[:, 0], trajectory[:, 1], trajectory[:, 2]
    dx, dy, dz = np.gradient(x)/dt, np.gradient(y)/dt, np.gradient(z)/dt
    speed = np.sqrt(dx**2 + dy**2 + dz**2) + 1e-8
    ux, uy, uz = dx/speed, dy/speed, dz/speed
    r = np.sqrt(x**2 + y**2) + 1e-8
    theta = np.arctan2(y, x)
    return np.stack([x, y, z, ux, uy, uz, r, np.sin(theta), np.cos(theta)], axis=-1)


def create_trajectory_windows(data, labels, context_len=60, pred_len=30, stride=15):
    """Create sliding-window samples from variable-length trajectories."""
    total_len = context_len + pred_len
    contexts, targets, sample_labels = [], [], []
    for i in range(len(data)):
        traj = data[i]
        valid_mask = ~np.isnan(traj[:, 0])
        valid_len = np.sum(valid_mask)
        if valid_len < total_len:
            continue
        traj_valid = traj[valid_mask]
        for start in range(0, valid_len - total_len + 1, stride):
            ctx_raw = traj_valid[start:start + context_len]
            tgt = traj_valid[start + context_len:start + total_len]
            ctx = compute_kinematic_features(ctx_raw)
            contexts.append(ctx)
            targets.append(tgt)
            sample_labels.append(labels[i])
    return (
        np.array(contexts, dtype=np.float32),
        np.array(targets, dtype=np.float32),
        np.array(sample_labels, dtype=np.int64),
    )


class AirTrackDataset(Dataset):
    """PyTorch Dataset for aircraft trajectory prediction."""
    def __init__(self, contexts, targets, labels):
        self.contexts = torch.from_numpy(contexts)
        self.targets = torch.from_numpy(targets)
        self.labels = torch.from_numpy(labels)
    
    def __len__(self):
        return len(self.contexts)
    
    def __getitem__(self, idx):
        return {"context": self.contexts[idx], "target": self.targets[idx], "label": self.labels[idx]}


def prepare_dataloaders(airport="RKSIa", context_len=60, pred_len=30, stride=15,
                        batch_size=32, cache_dir="./data/ATFMTraj", max_trajectories=None):
    """Full pipeline: download -> load -> window -> dataloader."""
    download_atfm_dataset(airport, cache_dir)
    train_data, train_labels = load_atfm_raw(airport, "TRAIN", cache_dir)
    test_data, test_labels = load_atfm_raw(airport, "TEST", cache_dir)
    if max_trajectories:
        train_data, train_labels = train_data[:max_trajectories], train_labels[:max_trajectories]
        test_data, test_labels = test_data[:max_trajectories], test_labels[:max_trajectories]
    
    train_ctx, train_tgt, train_lbl = create_trajectory_windows(train_data, train_labels, context_len, pred_len, stride)
    test_ctx, test_tgt, test_lbl = create_trajectory_windows(test_data, test_labels, context_len, pred_len, stride)
    
    all_labels = np.concatenate([train_lbl, test_lbl])
    n_classes = int(all_labels.max()) + 1
    
    train_ds = AirTrackDataset(train_ctx, train_tgt, train_lbl)
    test_ds = AirTrackDataset(test_ctx, test_tgt, test_lbl)
    train_loader = DataLoader(train_ds, batch_size=batch_size, shuffle=True, num_workers=2, pin_memory=True)
    test_loader = DataLoader(test_ds, batch_size=batch_size, shuffle=False, num_workers=2, pin_memory=True)
    
    return train_loader, test_loader, {
        "airport": airport, "context_len": context_len, "pred_len": pred_len,
        "n_features": train_ctx.shape[-1], "n_classes": n_classes,
        "n_train_windows": len(train_ds), "n_test_windows": len(test_ds),
    }