"""
SAIL Cryptocurrency DataLoader
================================
Parses historical cryptocurrency CSV files (like btc.csv, eth.csv) and converts
them into time-series forecasting and financial math training sequences.

This helps the AI develop numeric reasoning, pattern recognition in time series,
and understanding of financial data structures.
"""

import os
import torch
import pandas as pd
from torch.utils.data import Dataset, DataLoader

class CryptoDataset(Dataset):
    def __init__(self, data_dir, seq_length=128, max_files=None, tokenizer=None):
        """
        Args:
            data_dir: Directory containing the CSV files (e.g. datasets/finetune/archive)
            seq_length: Number of time steps to look back
            max_files: Max csv files to process (for debugging or quick runs)
            tokenizer: The BPE tokenizer to encode text representations
        """
        self.seq_length = seq_length
        self.tokenizer = tokenizer
        self.samples = []
        
        # Find all CSVs
        csv_files = [f for f in os.listdir(data_dir) if f.endswith('.csv')]
        if max_files:
            csv_files = csv_files[:max_files]
            
        print(f"Loading {len(csv_files)} cryptocurrency datasets...")
        
        for file in csv_files:
            coin_name = file.replace('.csv', '').upper()
            file_path = os.path.join(data_dir, file)
            
            try:
                # Read CSV - expected cols: time, id, amount, open, high, low, close, count, vol
                df = pd.read_csv(file_path)
                
                # Keep only numeric columns for sequence building
                # Typically we want open, high, low, close, vol
                cols = ['open', 'high', 'low', 'close', 'vol']
                
                # Check if columns exist
                if not all(c in df.columns for c in cols):
                    continue
                    
                df = df[cols].dropna()
                
                # We'll create string representations of the data to teach the model
                # financial math and forecasting in text format
                data_vals = df.values.tolist()
                
                # Create rolling windows of seq_length
                # Step by seq_length // 2 to get overlapping sequences but not explode RAM
                step = max(1, seq_length // 2)
                for i in range(0, len(data_vals) - seq_length, step):
                    window = data_vals[i : i + seq_length]
                    
                    # Create a prompt for the model
                    # For example, predicting the final closing price based on the series
                    series_str = []
                    for j, step_data in enumerate(window[:-1]): # Exclude target
                        o, h, l, c, v = step_data
                        series_str.append(f"T{j}: O={o:.2f} H={h:.2f} L={l:.2f} C={c:.2f} V={v:.2f}")
                        
                    context = " | ".join(series_str)
                    
                    target_step = window[-1]
                    target_c = target_step[3] # Close price
                    
                    instruction = f"Analyze the following {coin_name} price sequence and predict the next close price:\n{context}\n\nPrediction:"
                    response = f"Based on the sequence momentum and volatility, the predicted next close price is {target_c:.2f}."
                    
                    full_text = f"<|user|>\n{instruction}\n<|target|>\n{response}<|end|>"
                    self.samples.append(full_text)
                    
                    # Cap samples per coin to avoid massive RAM usage (some CSVs have 700k rows)
                    if len(self.samples) % 10000 == 0:
                        break # Only take first 10k windows per coin to balance dataset
                        
            except Exception as e:
                print(f"Error loading {file}: {e}")
                
        print(f"Loaded {len(self.samples)} crypto forecasting sequences.")

    def __len__(self):
        return len(self.samples)

    def __getitem__(self, idx):
        text = self.samples[idx]
        
        if self.tokenizer:
            tokens = self.tokenizer.encode(text)
            tokens = torch.tensor(tokens, dtype=torch.long)
            
            # Simple chunking for causal LM
            max_len = 1024
            if len(tokens) > max_len:
                tokens = tokens[:max_len]
                
            x = tokens[:-1]
            y = tokens[1:]
            
            # Pad if necessary
            if len(x) < max_len - 1:
                pad_len = max_len - 1 - len(x)
                x = torch.cat([x, torch.zeros(pad_len, dtype=torch.long)])
                y = torch.cat([y, torch.full((pad_len,), -100, dtype=torch.long)])
                
            return x, y
        
        return text

if __name__ == "__main__":
    # Test the loader
    loader = CryptoDataset("datasets/finetune/archive", seq_length=10, max_files=2)
    print(loader[0])