import jax
import json
import numpy as np
import optax
import model
import jax.numpy as jnp
from flax import nnx
from exp_utils import CFGDataLoader
from snapshot import Snapshot
from datetime import datetime
import mlflow
from util import load_reasoner
import logging
import sys

logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)


def loss_fn(instance, reasoner, data, target, mask, indices):
    embeddings = reasoner.embed(data)
    latent = reasoner.reason(embeddings)
    logits = instance(latent, input_tokens=data, latent_len=indices, token_len=indices)
    loss = optax.softmax_cross_entropy_with_integer_labels(logits=logits, labels=target) * mask

    return jnp.mean(loss)


@nnx.jit
def train(instance, reasoner, optimizer, metrics, data, target, mask, indices):
    grad_fn = nnx.value_and_grad(loss_fn)
    loss, grads = grad_fn(instance, reasoner, data, target, mask, indices)
    metrics.update(loss=loss)
    optimizer.update(instance, grads)


def main(reasoner_snap: str, size: str = 'large'):
    # Load configuration
    with open('configs/reasoner_cfg.json', 'r') as file:
        config = json.load(file)
        latent_feature = int(config[f'model_{size}']['Feature'])

    with open('configs/talker_cfg.json', 'r') as file:
        config = json.load(file)
        mlflow.log_params(config[f'model_{size}'])

        feature = int(config[f'model_{size}']['Feature'])
        attn_feature = int(config[f'model_{size}']['ATTN Feature'])
        ffn_feature = int(config[f'model_{size}']['FFN Feature'])
        num_head = int(config[f'model_{size}']['Head Count'])
        decoder_count = int(config[f'model_{size}']['Decoder Count'])
        encoder_count = int(config[f'model_{size}']['Encoder Count'])
        init_scalar = float(config[f'model_{size}']['Init Scalar'])
        max_len = int(config[f'model_{size}']['Max Length'])
        rope_base = float(config[f'model_{size}']['RoPE Base'])

        mlflow.log_params(config['adaptive'])

        context_len = int(config['adaptive']['Context Length'])
        peak_lr = float(config['adaptive']['Peak LR'])
        gradient_clipping = float(config['adaptive']['Grads Clipping'])
        weight_decay = float(config['adaptive']['Weight Decay'])
        total_steps = int(config['adaptive']['Total Steps'])
        warmup_steps = int(config['adaptive']['Warmup Steps'])
        anneal_steps = int(config['adaptive']['Anneal Steps'])
        batch_size = int(config['adaptive']['Batch Size'])

        if 'Accumulation' in config['adaptive']:
            accu_steps = int(config['adaptive']['Accumulation'])
        else:
            accu_steps = 1

        if size != 'large':
            accu_steps /= 2
            accu_steps = int(accu_steps)
            batch_size *= 2

    # Build model instance
    key = jax.random.key(0)
    reasoner = load_reasoner(reasoner_snap, vocab_size=32, size=size)
    instance = model.DualTalker(
        feature=feature,
        latent_feature=latent_feature,
        attn_feature=attn_feature,
        ffn_feature=ffn_feature,
        num_head=num_head,
        encoder_count=encoder_count,
        decoder_count=decoder_count,
        init_scalar=init_scalar,
        vocab_size=32,
        key=key,
        dtype=jnp.bfloat16
    )
    instance.train(
        rope_base=rope_base,
        max_len=max_len
    )

    # Build optimizer & learning rate scheduler
    scheduler = optax.join_schedules(
        schedules=[
            optax.linear_schedule(
                init_value=0.0,
                end_value=peak_lr,
                transition_steps=warmup_steps,
            ),
            optax.constant_schedule(value=peak_lr),
            optax.linear_schedule(
                init_value=peak_lr,
                end_value=0.0,
                transition_steps=anneal_steps
            )
        ],
        boundaries=[warmup_steps, total_steps - anneal_steps]
    )

    optimizer = optax.chain(
        optax.clip_by_global_norm(gradient_clipping),
        optax.adamw(
            learning_rate=scheduler,
            weight_decay=weight_decay,
            b1=0.9,
            b2=0.95,
            eps=1e-6
        )
    )

    optimizer = nnx.Optimizer(instance, optax.MultiSteps(optimizer, every_k_schedule=accu_steps), wrt=nnx.Param)
    metrics = nnx.MultiMetric(loss=nnx.metrics.Average('loss'))
    snap = Snapshot(path=datetime.now().strftime(f'snapshot/cfg_completion/talker/adaptive_cfg3f_{size}_%d-%m-%y_%H:%M'))

    # Data loader setup
    train_loader = CFGDataLoader(
        path=[PATH_TO_CFG_DATA],
        pattern='cfg_train*.parquet',
        tokenizer='tokenizer_250725',
        pad_token=0,
        mem=128,
        batch_size=batch_size,
        context_len=context_len,
        threads=1,
        dtype=np.int32,
    )

    # Create logger
    def log_and_save():
        train_loss = metrics.compute().get('loss')
        logger.info(train_loss)
        grads = optimizer.opt_state.acc_grads
        metrics.reset()

        try:
            mlflow.log_metrics({
                'train loss': train_loss,
                'grad: global': optax.global_norm(grads),
                'grad: embed': optax.global_norm(grads['featurizer']),
                'grad: final norm': optax.global_norm(grads['norm']),
                'learning rate': scheduler(step // accu_steps),
            }, step=step // accu_steps)

            for i in range(encoder_count):
                mlflow.log_metrics({
                    f'grad: encoder-{i:02d}': float(optax.global_norm(grads['encoders'][i])),
                    f'grad: attention-{i:02d}': float(optax.global_norm(grads['encoders'][i]['attention'])),
                }, step=step // accu_steps)

            for i in range(decoder_count):
                mlflow.log_metrics({
                    f'grad: decoder-{i:02d}': float(optax.global_norm(grads['decoders'][i])),
                    f'grad: self-attention-{i:02d}': float(optax.global_norm(grads['decoders'][i]['self_attention']['w_k'])),
                    f'grad: cross-attention-{i:02d}': float(optax.global_norm(grads['decoders'][i]['cross_attention']['w_k'])),
                }, step=step // accu_steps)

        except Exception as e:
            print(f'Logging error at step {step}: {e}\n')

        snap.save(f'cfg3f_{size}_{step // accu_steps}', instance)

    # Train
    for step in range(total_steps * accu_steps):
        data, target, mask, indices = next(train_loader)
        train(instance, reasoner, optimizer, metrics, data, target, mask, indices)

        if step % (accu_steps * 100) == 0:
            log_and_save()

    snap.save(f'cfg3f_{size}_{total_steps}', instance)
    train_loader.stop()


if __name__ == '__main__':
    from util import prepare_training_env
    PATH_TO_CFG_DATA = './data'
    model_size = sys.argv[1]
    reasoner_snap = sys.argv[2]

    mlflow_credential = 'YOUR_MLFLOW_CREDENTIAL'
    mlflow_tracking_uri = 'YOUR_MLFLOW_TRACKING_URI'
    gpu_preallocation_fraction = '.99'
    prepare_training_env(mlflow_tracking_uri, gpu_preallocation_fraction, mlflow_credential)
    mlflow.set_experiment('Talker')

    with mlflow.start_run(run_name=datetime.now().strftime('train_%d/%m/%y-%H:%M')):
        main(reasoner_snap, size=model_size)
