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-025000-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 fn_query_on_load():
    return "Biofuels would disrupt"


def generate_output(prompt, max_new_tokens=200, temperature=0.8, top_k=50):
    m.eval()
    encoded_text = tokenizer_gpt.encode(prompt)
    # print('--------------------encoded text = ',encoded_text)

    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=max_new_tokens, temperature=0.8, top_k=50)[0])

    m.train()
    return {
        output: out_text
    }


with gr.Blocks() as app:
    with gr.Row():
        gr.Markdown(
            """
            # MiniGPT - GPT Training on LLaMa with redpajama dataset
            ###  Enter a context to generate automated text "
            """)

    with gr.Row(visible=True):
        search_text = gr.Textbox(value=fn_query_on_load, placeholder='Enter prompt..', label='Enter Prompt')

    with gr.Row():
        submit_btn = gr.Button("Submit", variant='primary')
        clear_btn = gr.ClearButton()
    with gr.Row():
        with gr.Row():
            output = gr.Textbox(lines=15, interactive=False, label='Out Box')

    def clear_data():
        return {
            output: None,
            search_text: None
        }

    clear_btn.click(clear_data, None, [output, search_text])


    submit_btn.click(
        generate_output,
        search_text,
        output
    )


'''
Launch the app
'''
app.queue().launch()