# train_rwkv_flax.py

import jax
import jax.numpy as jnp
from flax.training import train_state
from flax import linen as nn
import optax
import wandb
from transformers import AutoTokenizer
from datasets import load_dataset
from dataclasses import dataclass
import numpy as np
import os
from rwkv_flax import RWKV, Config  # Ensure rwkv_flax.py is in the same directory or adjust the import accordingly

# Define the training state using Flax's train_state
@dataclass
class TrainStateWithBatchStats(train_state.TrainState):
    batch_stats: Any = None  # Modify as needed

# Function to create the training state
def create_train_state(rng, model, learning_rate):
    params = model.init(rng, jnp.ones((1, 2), dtype=jnp.int32))['params']
    tx = optax.adam(learning_rate)
    return train_state.TrainState.create(
        apply_fn=model.apply,
        params=params,
        tx=tx
    )

# Define the loss function (Cross-Entropy Loss)
def compute_loss(logits, labels):
    one_hot = jax.nn.one_hot(labels, logits.shape[-1])
    loss = optax.softmax_cross_entropy(logits=logits, labels=one_hot)
    return loss.mean()

# Define the training step
@jax.jit
def train_step(state, batch):
    def loss_fn(params):
        logits = state.apply_fn({'params': params}, batch['input_ids'])
        loss = compute_loss(logits, batch['labels'])
        return loss, logits

    grad_fn = jax.value_and_grad(loss_fn, has_aux=True)
    (loss, logits), grads = grad_fn(state.params)
    state = state.apply_gradients(grads=grads)
    return state, loss

# Load dataset and tokenizer
def get_dataset(tokenizer, block_size=128):
    base_url = "https://huggingface.co/datasets/chengjunyan1/smollm-12.5-corpus/resolve/main/fineweb-edu-dedup/"

    data_files = {"train": base_url+"train-0000[1]-of-00237.parquet"} #,2,3,4,5,6
    
    ds = load_dataset("parquet", data_files=data_files, keep_in_memory=True)
    ds = ds.select_columns(['text'])
    dataset = ds["train"]

    def tokenize_function(examples):
        return tokenizer(examples['text'])

    tokenized = dataset.map(
        tokenize_function, 
        batched=True, 
        remove_columns=['text']
    )

    def group_texts(examples):
        concatenated = {k: sum(examples[k], []) for k in examples.keys()}
        total_length = len(concatenated[list(concatenated.keys())[0]])
        total_length = (total_length // block_size) * block_size
        result = {
            k: [t[i : i + block_size] for i in range(0, total_length, block_size)]
            for k, t in concatenated.items()
        }
        result['labels'] = result['input_ids'].copy()
        return result

    lm_dataset = tokenized.map(
        group_texts,
        batched=True,
        batch_size=128,
        #num_proc=60,
    )

    lm_dataset = lm_dataset.with_format('jax', columns=['input_ids', 'labels'])

    return lm_dataset

# Initialize wandb
wandb.init(
    project="RWKV-Flax-Training",
    config=Config(),
)
config = wandb.config

# Initialize tokenizer
tokenizer = AutoTokenizer.from_pretrained("gpt2")  # Replace with appropriate tokenizer

# Load dataset
dataset = get_dataset(tokenizer)
train_loader = jax.device_put(dataset)

# Initialize model
rng = jax.random.PRNGKey(0)
model = RWKV(cfg=config)
train_rng, init_rng = jax.random.split(rng)
state = create_train_state(init_rng, model, learning_rate=1e-4)

# Training loop parameters
num_epochs = 3
batch_size = 32
steps_per_epoch = len(dataset) // batch_size

# Start training
for epoch in range(num_epochs):
    epoch_loss = 0.0
    for step, batch in enumerate(train_loader.shard(jax.local_device_count(), index=jax.process_index())):
        state, loss = train_step(state, batch)
        epoch_loss += loss
        if step % 100 == 0:
            print(f"Epoch {epoch}, Step {step}, Loss: {loss}")
    avg_loss = epoch_loss / steps_per_epoch
    wandb.log({"epoch": epoch, "loss": avg_loss})
    print(f"Epoch {epoch} completed. Average Loss: {avg_loss}")

# Save the trained model
if jax.process_index() == 0:
    model_dir = "rwkv_flax_model"
    os.makedirs(model_dir, exist_ok=True)
    np.savez(os.path.join(model_dir, "params.npz"), **state.params)
    wandb.save(os.path.join(model_dir, "params.npz"))

wandb.finish()
use half precision