import torch
import torch.nn as nn
import gradio as gr
from tsai_gpt.tokenizer import Tokenizer
import lightning as L
from lightning.fabric.loggers import CSVLogger
from pathlib import Path
from tsai_gpt.utils import  num_parameters, load_checkpoint, get_default_supported_precision
from tsai_gpt.model import GPT, Block, Config

model_name = "pythia-160m"
name = "redpajama"
out_dir = Path("out") / name
log_interval = 100

precision = get_default_supported_precision(False)
logger = CSVLogger("out", name, flush_logs_every_n_steps=log_interval)
fabric = L.Fabric(devices=1, strategy="auto", precision=precision, loggers=logger)

config = Config.from_name(model_name)

def _init_weights(module: nn.Module) -> None:
        """Meant to be used with `gpt.apply(gpt._init_weights)`."""
        if isinstance(module, nn.Linear):
            torch.nn.init.normal_(module.weight, mean=0.0, std=0.02)
            if module.bias is not None:
                torch.nn.init.zeros_(module.bias)
        elif isinstance(module, nn.Embedding):
            torch.nn.init.normal_(module.weight, mean=0.0, std=0.02)
        
with fabric.init_module(empty_init=True):
    model = GPT(config)
    model.apply(_init_weights)
model.apply(_init_weights)


checkpoint_path = Path("out/redpajama/iter-015000-ckpt.pth")

load_checkpoint(fabric, model, checkpoint_path)
    
#print(model.transformer.h[0].mlp.fc.weight)

#fabric.print(f"Time to instantiate model: {time.perf_counter() - t0:.02f} seconds.")
#fabric.print(f"Total parameters {num_parameters(model):,}")

weight_decay = 1e-1
beta1 = 0.9
beta2 = 0.95
learning_rate = 6e-3
hparams = {k: v for k, v in locals().items() if isinstance(v, (int, float, str)) and not k.startswith("_")}

model = fabric.setup(model)
optimizer = torch.optim.AdamW(
    model.parameters(), lr=learning_rate, weight_decay=weight_decay, betas=(beta1, beta2), foreach=False
)

# model_copy = model

optimizer = fabric.setup_optimizers(optimizer)

state = {"model": model, "optimizer": optimizer, "hparams": hparams, "iter_num": 0, "step_count": 0}

resume = max(out_dir.glob("*.pth"), key=lambda p: int(p.name.split("-")[1]))
if resume:
    fabric.print(f"Loading model from {resume}")
    fabric.load(resume, state)

deviceType = 'cuda' if torch.cuda.is_available() else 'cpu'
m  = model.to(deviceType)
tokenizer_gpt = Tokenizer(checkpoint_dir=Path("checkpoints/meta-llama/Llama-2-7b-chat-hf")) 


def inference(input_context, count):
    #print('--------------------input = ',input_context)
    encoded_text = tokenizer_gpt.encode(input_context)
    #print('--------------------encoded text = ',encoded_text)
    count = int(count)
    #print('--------------------count = ',count)
    reshaped_tensor = torch.unsqueeze(encoded_text, 0).to(deviceType)     
    #print('--------------------reshaped_tensor = ',reshaped_tensor)
    out_text = tokenizer_gpt.decode(m.generate(reshaped_tensor, max_new_tokens=count)[0])
    return out_text

title = "TSAI S22 Assignment: GPT training on LLaMa - redpajama dataset"
description = "A simple Gradio interface that accepts a context and generates text "
examples = [["Machine Learning","200"],
            ["Deep Learning","200"]
           ]
 

demo = gr.Interface(
    inference, 
    inputs = [gr.Textbox(placeholder="Enter starting characters"), gr.Textbox(placeholder="Enter number of characters you want to generate")], 
    outputs = [gr.Textbox(label="Generated text")],
    title = title,
    description = description,
    examples = examples
)
demo.launch()