from transformers import LlamaConfig, AutoModelForCausalLM, AutoTokenizer, Trainer, TrainingArguments
from datasets import Dataset
import pandas as pd
import numpy as np
import torch
from pathlib import Path

def generate_test_candles(n_samples=50):
    timestamps = pd.date_range(start="2024-01-01", periods=n_samples, freq="D")
    data = {
        "time": timestamps.astype(np.int64) // 10**9,
        "open": np.random.uniform(100, 150, n_samples),
        "high": np.random.uniform(150, 200, n_samples),
        "low": np.random.uniform(50, 100, n_samples),
        "close": np.random.uniform(100, 150, n_samples),
        "volume": np.random.uniform(1000, 5000, n_samples)
    }
    return pd.DataFrame(data)

def prepare_training_data(examples):
    features = np.column_stack([
        examples['time'],
        examples['open'],
        examples['high'],
        examples['low'],
        examples['close'],
        examples['volume']
    ])
    
    return {
        "input_ids": torch.tensor(features, dtype=torch.float32),
        "labels": torch.tensor(examples['close'], dtype=torch.float32)
    }

Path("./results").mkdir(exist_ok=True)
Path("./logs").mkdir(exist_ok=True)

test_data = generate_test_candles()
dataset = Dataset.from_pandas(test_data)
dataset = dataset.map(prepare_training_data, batched=True)

training_args = TrainingArguments(
    output_dir="./results",
    num_train_epochs=3,
    per_device_train_batch_size=4,
    learning_rate=2e-4,
    fp16=True,
    save_steps=100,
    logging_dir="./logs",
    logging_steps=10,
    remove_unused_columns=False
)

config = LlamaConfig()
model = AutoModelForCausalLM.from_pretrained(
    "tmmdev2/llama-vision-charts",
    config=config,
    trust_remote_code=True
)

tokenizer = AutoTokenizer.from_pretrained(
    "tmmdev2/llama-vision-charts",
    trust_remote_code=True
)

trainer = Trainer(
    model=model,
    args=training_args,
    train_dataset=dataset,
    tokenizer=tokenizer
)

results = trainer.train()
model.save_pretrained("./results/final_model")
tokenizer.save_pretrained("./results/final_model")

print("Training completed!")
print(f"Final loss: {results.training_loss}")