"""
predictor.py — Kronos inference that retains ALL sampled paths.

Same as the local version but imports Kronos from the `Kronos/` subfolder
that the Space's pre_build step clones from GitHub.
"""
from __future__ import annotations

import os
import sys
import logging
from pathlib import Path

import numpy as np
import pandas as pd
import torch

# Kronos model code lives at ./model/ in this Space (copied from upstream)
SPACE_ROOT = Path(__file__).resolve().parent
if str(SPACE_ROOT) not in sys.path:
    sys.path.insert(0, str(SPACE_ROOT))

from model import Kronos, KronosTokenizer, KronosPredictor  # noqa: E402
from model.kronos import sample_from_logits, calc_time_stamps  # noqa: E402

logger = logging.getLogger(__name__)

MODEL_IDS = {
    "mini": ("NeoQuasar/Kronos-Tokenizer-2k", "NeoQuasar/Kronos-mini", 2048),
    "small": ("NeoQuasar/Kronos-Tokenizer-base", "NeoQuasar/Kronos-small", 512),
    "base": ("NeoQuasar/Kronos-Tokenizer-base", "NeoQuasar/Kronos-base", 512),
}


def _auto_regressive_inference_keep_paths(
    tokenizer, model,
    x, x_stamp, y_stamp,
    max_context: int, pred_len: int,
    clip: float = 5.0,
    T: float = 1.0,
    top_k: int = 0,
    top_p: float = 0.9,
    sample_count: int = 5,
    verbose: bool = False,
):
    """Copy of Kronos's auto_regressive_inference with averaging removed."""
    try:
        from tqdm import trange
    except ImportError:
        trange = range

    with torch.no_grad():
        x = torch.clip(x, -clip, clip)
        device = x.device

        x = x.unsqueeze(1).repeat(1, sample_count, 1, 1).reshape(-1, x.size(1), x.size(2)).to(device)
        x_stamp = x_stamp.unsqueeze(1).repeat(1, sample_count, 1, 1).reshape(-1, x_stamp.size(1), x_stamp.size(2)).to(device)
        y_stamp = y_stamp.unsqueeze(1).repeat(1, sample_count, 1, 1).reshape(-1, y_stamp.size(1), y_stamp.size(2)).to(device)

        x_token = tokenizer.encode(x, half=True)

        initial_seq_len = x.size(1)
        batch_size = x_token[0].size(0)
        total_seq_len = initial_seq_len + pred_len
        full_stamp = torch.cat([x_stamp, y_stamp], dim=1)

        generated_pre = x_token[0].new_empty(batch_size, pred_len)
        generated_post = x_token[1].new_empty(batch_size, pred_len)

        pre_buffer = x_token[0].new_zeros(batch_size, max_context)
        post_buffer = x_token[1].new_zeros(batch_size, max_context)
        buffer_len = min(initial_seq_len, max_context)
        if buffer_len > 0:
            start_idx = max(0, initial_seq_len - max_context)
            pre_buffer[:, :buffer_len] = x_token[0][:, start_idx:start_idx + buffer_len]
            post_buffer[:, :buffer_len] = x_token[1][:, start_idx:start_idx + buffer_len]

        ran = trange if verbose else range
        for i in ran(pred_len):
            current_seq_len = initial_seq_len + i
            window_len = min(current_seq_len, max_context)

            if current_seq_len <= max_context:
                input_tokens = [
                    pre_buffer[:, :window_len],
                    post_buffer[:, :window_len],
                ]
            else:
                input_tokens = [pre_buffer, post_buffer]

            context_end = current_seq_len
            context_start = max(0, context_end - max_context)
            current_stamp = full_stamp[:, context_start:context_end, :].contiguous()

            s1_logits, context = model.decode_s1(input_tokens[0], input_tokens[1], current_stamp)
            s1_logits = s1_logits[:, -1, :]
            sample_pre = sample_from_logits(s1_logits, temperature=T, top_k=top_k, top_p=top_p, sample_logits=True)

            s2_logits = model.decode_s2(context, sample_pre)
            s2_logits = s2_logits[:, -1, :]
            sample_post = sample_from_logits(s2_logits, temperature=T, top_k=top_k, top_p=top_p, sample_logits=True)

            generated_pre[:, i] = sample_pre.squeeze(-1)
            generated_post[:, i] = sample_post.squeeze(-1)

            if current_seq_len < max_context:
                pre_buffer[:, current_seq_len] = sample_pre.squeeze(-1)
                post_buffer[:, current_seq_len] = sample_post.squeeze(-1)
            else:
                pre_buffer.copy_(torch.roll(pre_buffer, shifts=-1, dims=1))
                post_buffer.copy_(torch.roll(post_buffer, shifts=-1, dims=1))
                pre_buffer[:, -1] = sample_pre.squeeze(-1)
                post_buffer[:, -1] = sample_post.squeeze(-1)

        full_pre = torch.cat([x_token[0], generated_pre], dim=1)
        full_post = torch.cat([x_token[1], generated_post], dim=1)

        context_start = max(0, total_seq_len - max_context)
        input_tokens = [
            full_pre[:, context_start:total_seq_len].contiguous(),
            full_post[:, context_start:total_seq_len].contiguous(),
        ]
        z = tokenizer.decode(input_tokens, half=True)
        z = z.reshape(-1, sample_count, z.size(1), z.size(2))
        # NOT averaging; keep all paths
        preds = z.cpu().numpy()
        return preds


class KronosPathsPredictor:
    def __init__(self, model_size: str = "small", device: str | None = None):
        if model_size not in MODEL_IDS:
            raise ValueError(f"model_size must be one of {list(MODEL_IDS.keys())}")
        tokenizer_id, model_id, max_context = MODEL_IDS[model_size]
        self.max_context = max_context
        self.model_size = model_size
        self.model_id = model_id

        if device is None:
            if torch.cuda.is_available():
                device = "cuda:0"
            elif hasattr(torch.backends, "mps") and torch.backends.mps.is_available():
                device = "mps"
            else:
                device = "cpu"
        self.device = device

        logger.info(f"Loading tokenizer: {tokenizer_id}")
        self.tokenizer = KronosTokenizer.from_pretrained(tokenizer_id)
        logger.info(f"Loading model: {model_id}")
        self.model = Kronos.from_pretrained(model_id)
        self.tokenizer = self.tokenizer.to(self.device)
        self.model = self.model.to(self.device)

        self.price_cols = ["open", "high", "low", "close"]
        self.vol_col = "volume"
        self.amt_col = "amount"
        self.clip = 5.0

    def _prepare_one(self, df: pd.DataFrame, x_timestamp: pd.Series, y_timestamp: pd.Series):
        df = df.copy()
        if self.vol_col not in df.columns:
            df[self.vol_col] = 0.0
            df[self.amt_col] = 0.0
        if self.amt_col not in df.columns:
            df[self.amt_col] = df[self.vol_col] * df[self.price_cols].mean(axis=1)

        if df[self.price_cols + [self.vol_col, self.amt_col]].isnull().values.any():
            raise ValueError("DataFrame contains NaN in price/volume columns")

        x_time = calc_time_stamps(x_timestamp)
        y_time = calc_time_stamps(y_timestamp)

        x = df[self.price_cols + [self.vol_col, self.amt_col]].values.astype(np.float32)
        x_stamp = x_time.values.astype(np.float32)
        y_stamp = y_time.values.astype(np.float32)

        x_mean = np.mean(x, axis=0)
        x_std = np.std(x, axis=0)
        x_norm = (x - x_mean) / (x_std + 1e-5)
        x_norm = np.clip(x_norm, -self.clip, self.clip)

        return x_norm, x_stamp, y_stamp, x_mean, x_std

    def predict_paths(
        self,
        df_list, x_timestamp_list, y_timestamp_list,
        pred_len: int,
        T: float = 1.0, top_k: int = 0, top_p: float = 0.9,
        sample_count: int = 30,
        verbose: bool = False,
        seed: int | None = None,
    ):
        if seed is not None:
            torch.manual_seed(seed)
            np.random.seed(seed)

        if not (len(df_list) == len(x_timestamp_list) == len(y_timestamp_list)):
            raise ValueError("df_list / x_timestamp_list / y_timestamp_list must be same length")

        seq_lens = [len(df) for df in df_list]
        y_lens = [len(ts) for ts in y_timestamp_list]
        if len(set(seq_lens)) > 1:
            raise ValueError(f"All dfs must have same historical length, got {seq_lens}")
        if len(set(y_lens)) > 1:
            raise ValueError(f"All y_timestamps must have same length, got {y_lens}")

        x_batch, x_stamp_batch, y_stamp_batch, means, stds = [], [], [], [], []
        for df, x_ts, y_ts in zip(df_list, x_timestamp_list, y_timestamp_list):
            x_norm, x_stamp, y_stamp, m, s = self._prepare_one(df, x_ts, y_ts)
            x_batch.append(x_norm); x_stamp_batch.append(x_stamp); y_stamp_batch.append(y_stamp)
            means.append(m); stds.append(s)

        x_batch = np.stack(x_batch, axis=0)
        x_stamp_batch = np.stack(x_stamp_batch, axis=0)
        y_stamp_batch = np.stack(y_stamp_batch, axis=0)

        x_tensor = torch.from_numpy(x_batch.astype(np.float32)).to(self.device)
        x_stamp_tensor = torch.from_numpy(x_stamp_batch.astype(np.float32)).to(self.device)
        y_stamp_tensor = torch.from_numpy(y_stamp_batch.astype(np.float32)).to(self.device)

        preds_all = _auto_regressive_inference_keep_paths(
            self.tokenizer, self.model,
            x_tensor, x_stamp_tensor, y_stamp_tensor,
            self.max_context, pred_len,
            self.clip, T, top_k, top_p, sample_count, verbose,
        )

        results = []
        for i, (m, s) in enumerate(zip(means, stds)):
            pred = preds_all[i]
            pred = pred[:, -pred_len:, :]
            pred = pred * (s + 1e-5) + m
            results.append(pred.astype(np.float32))
        return results