import torch
from transformers import AutoModelForCausalLM, AutoTokenizer, AutoConfig, TrainingArguments, Trainer, BertLMHeadModel, BertForSequenceClassification
from datasets import Dataset
import pandas as pd
import csv
from transformers import TrainingArguments, Trainer

import tensorflow as tf
# Check TensorFlow GPU availability
print("GPUs Available: ", tf.config.list_physical_devices('GPU'))

import os
# Setting the environment variable for MPS
os.environ['PYTORCH_MPS_HIGH_WATERMARK_RATIO'] = '0.0'

def get_device():
    """Automatically chooses the best device."""
    if torch.cuda.is_available():
        return torch.device('cuda')
    elif torch.backends.mps.is_available():
        return torch.device('mps')
    else:
        return torch.device('cpu')

def load_data_and_config(data_path):
    """Loads training data from CSV."""
    data = []
    with open(data_path, newline='', encoding='utf-8') as csvfile:
        reader = csv.DictReader(csvfile, delimiter=';')
        for row in reader:
            data.append({'text': row['description']})
    return data

def train_model(model, tokenizer, data, device):
    """Trains the model using the Hugging Face Trainer API."""
    # inputs = [tokenizer(d['text'], max_length=256, truncation=True, padding='max_length', return_tensors="pt") for d in data]
    inputs = [tokenizer(d['text'], max_length=256, truncation=True, padding='max_length', return_tensors="pt").to(torch.float16) for d in data]
    dataset = Dataset.from_dict({
        'input_ids': [x['input_ids'].squeeze() for x in inputs],
        'labels': [x['input_ids'].squeeze() for x in inputs]
    })
    
    training_args = TrainingArguments(
        output_dir='./results',
        num_train_epochs=3,
        per_device_train_batch_size=1,
        gradient_accumulation_steps=4,
        fp16=True,  # Enable mixed precision
        warmup_steps=500,
        weight_decay=0.01,
        logging_dir='./logs',
        logging_steps=10,
    )
    
    trainer = Trainer(
        model=model,
        args=training_args,
        train_dataset=dataset,
        tokenizer=tokenizer
    )
    
    trainer.train()

     # Optionally clear cache if using GPU or MPS
    if torch.cuda.is_available():
        print(torch.cuda.memory_summary(device=None, abbreviated=False))
        torch.cuda.empty_cache()
    elif torch.has_mps:
        torch.mps.empty_cache()

    # Perform any remaining steps such as logging, saving, etc.
    trainer.save_model()

def main(api_name, base_url):
    device = get_device()  # Get the appropriate device
    data = load_data_and_config("train2.csv")
    tokenizer = AutoTokenizer.from_pretrained("google/codegemma-2b")
    # Load the configuration for a specific model
    config = AutoConfig.from_pretrained('google/codegemma-2b')
    # Update the activation function
    # config.hidden_act = ''  # Set to use approximate GeLU gelu_pytorch_tanh
    config.hidden_activation = 'gelu_pytorch_tanh'  # Set to use GeLU

    model = AutoModelForCausalLM.from_pretrained('google/codegemma-2b', is_decoder=True)
    #model = BertLMHeadModel.from_pretrained('google/codegemma-2b', is_decoder=True)
    # Example assuming you have a prepared dataset for classification
    #model = BertForSequenceClassification.from_pretrained('thenlper/gte-small', num_labels=2, is_decoder=True)  # binary classification
    
    model.to(device)  # Move model to the appropriate device

    train_model(model, tokenizer, data, device)
    
    model.save_pretrained("./fine_tuned_model")
    tokenizer.save_pretrained("./fine_tuned_model")

    prompt = "I need to retrieve the latest block on chain using a python script"
    api_query = generate_api_query(model, tokenizer, prompt, "latest block on chain", api_name, base_url)
    print(f"Generated code: {api_query}")

def generate_api_query(model, tokenizer, prompt, desired_output, api_name, base_url):
    # Prepare input prompt for the model, ensure tensors are compatible with PyTorch
    input_ids = tokenizer.encode(f"{prompt} Write an API query to {api_name} to get {desired_output}", return_tensors="pt")

    # Ensure input_ids are on the same device as the model
    input_ids = input_ids.to(model.device)

    # Generate query using model with temperature for randomness
    output = model.generate(input_ids, max_length=128, temperature=0.001, do_sample=True)
    

    # Decode the generated query tokens
    query = tokenizer.decode(output[0], skip_special_tokens=True)
    return f"{base_url}/{query}"

if __name__ == "__main__":
    api_name = "Koios"
    base_url = "https://api.koios.rest/v1"
    main(api_name, base_url)