import json
import itertools
import torch
from torch.utils.data import Dataset, DataLoader
from transformers import LlamaTokenizer, LlamaForCausalLM, Trainer, TrainingArguments

# Load JSON files
def load_json(file_path):
    with open(file_path, 'r') as file:
        return json.load(file)

# Generate prompts based on schema and data
def generate_prompts(schema, data):
    prompts = []

    for table_name, columns in schema.items():
        table_data = data.get(table_name, [])
        
        # Generate all possible column combinations
        column_combinations = []
        for r in range(1, len(columns) + 1):
            column_combinations.extend(itertools.combinations(columns, r))
        
        # Generate prompts for each row in data
        for row in table_data:
            for combination in column_combinations:
                prompt_parts = []
                for column in combination:
                    if column in row:
                        prompt_parts.append(f"{column}: {row[column]}")
                prompt = f"Table: {table_name}\n" + "\n".join(prompt_parts)
                prompts.append(prompt)
    
    return prompts

# Dataset class for the prompts
class PromptsDataset(Dataset):
    def __init__(self, prompts, tokenizer, max_length):
        self.prompts = prompts
        self.tokenizer = tokenizer
        self.max_length = max_length

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

    def __getitem__(self, idx):
        prompt = self.prompts[idx]
        encoding = self.tokenizer(prompt, return_tensors='pt', truncation=True, padding='max_length', max_length=self.max_length)
        input_ids = encoding['input_ids'].squeeze()
        attention_mask = encoding['attention_mask'].squeeze()
        return {'input_ids': input_ids, 'attention_mask': attention_mask, 'labels': input_ids}

# Main function
def main():
    schema_file_path = 'schema7878.json'
    data_file_path = 'data.json'
    output_file_path = 'prompts.txt'

    schema = load_json(schema_file_path)
    data = load_json(data_file_path)
    
    prompts = generate_prompts(schema, data)
    
    # Save prompts to a file (optional)
    with open(output_file_path, 'w') as file:
        for prompt in prompts:
            file.write(prompt + "\n\n")

    # Load tokenizer and model for Llama 3
    tokenizer = LlamaTokenizer.from_pretrained('llama-3')
    model = LlamaForCausalLM.from_pretrained('llama-3')

    # Create dataset and dataloader
    dataset = PromptsDataset(prompts, tokenizer, max_length=512)
    dataloader = DataLoader(dataset, batch_size=4, shuffle=True)

    # Define training arguments
    training_args = TrainingArguments(
        output_dir='./results',
        overwrite_output_dir=True,
        num_train_epochs=3,
        per_device_train_batch_size=4,
        save_steps=10_000,
        save_total_limit=2,
        prediction_loss_only=True,
    )

    # Initialize Trainer
    trainer = Trainer(
        model=model,
        args=training_args,
        train_dataset=dataset,
    )

    # Train the model
    trainer.train()

    # Save the trained model
    model.save_pretrained('./trained_model')
    tokenizer.save_pretrained('./trained_model')
    print(f"Generated {len(prompts)} prompts and saved to {output_file_path}")

if __name__ == "__main__":
    main()