"""
This is an experiment to:
* entertain a thought to try and abide by HF's transformers API (to benefit from caching better)
* conform to a single embedding (instead of a bunch of them) by folding/unfolding inputs
* stop trying to make a mixed AR+NAR model work since it seems lobotomized if I keep trying to enforce both recurrent and parallel inferencing (despite a penalty cost)
	+ I will not cave and go with codebook patterns, not yet.
"""

from ..config import cfg

from ..data import fold_inputs, unfold_outputs

import torch
import torch.nn.functional as F
from torch.nn.utils.rnn import pad_sequence
from torch import Tensor
from torch.nn import CrossEntropyLoss
from torch.utils.checkpoint import checkpoint
from torchmetrics.classification import BinaryAccuracy, MulticlassAccuracy, MulticlassPrecision

import random
import math

from einops import rearrange
from tqdm import trange

AVAILABLE_ARCHES = []

try:
	from transformers import LlamaForCausalLM, LlamaConfig
	AVAILABLE_ARCHES.append("llama")
except Exception as e:
	print("Error importing `llama` arch:", e)
	pass

try:
	from .retnet_hf import RetNetConfig
	from ..ext.retnet_hf.modeling_retnet import RetNetForCausalLM

	AVAILABLE_ARCHES.append("retnet")
except Exception as e:
	print("Error importing `retnet` arch:", e)
	pass

try:
	from mamba_ssm.models.mixer_seq_simple import MambaLMHeadModel, MambaConfig, MixerModel as MambaMixelModel, layer_norm_fn as MambaLayerNormFn, RMSNorm as MambaRMSNorm

	def MambaMixelModel_forward(self, input_ids, inference_params=None, **mixer_kwargs):
		hidden_states = self.embedding(input_ids)
		residual = None
		for layer in self.layers:
			if self.gradient_checkpointing and hidden_states.requires_grad:
				hidden_states, residual = checkpoint( layer, hidden_states, residual, inference_params=inference_params, use_reentrant=False )
			else:
				hidden_states, residual = layer( hidden_states, residual, inference_params=inference_params )
		if not self.fused_add_norm:
			residual = (hidden_states + residual) if residual is not None else hidden_states
			hidden_states = self.norm_f(residual.to(dtype=self.norm_f.weight.dtype))
		else:
			# Set prenorm=False here since we don't need the residual
			hidden_states = MambaLayerNormFn(
				hidden_states,
				self.norm_f.weight,
				self.norm_f.bias,
				eps=self.norm_f.eps,
				residual=residual,
				prenorm=False,
				residual_in_fp32=self.residual_in_fp32,
				is_rms_norm=isinstance(self.norm_f, MambaRMSNorm)
			)
		return hidden_states

	MambaMixelModel.forward = MambaMixelModel_forward

	AVAILABLE_ARCHES.append("mamba")
	AVAILABLE_ARCHES.append("mamba2")
except Exception as e:
	print("Error importing `mamba` arch:", e)
	pass


SELECTED_ARCH = cfg.model.arch_type 
if SELECTED_ARCH not in AVAILABLE_ARCHES:
	raise ValueError(f"Requesting arch `{SELECTED_ARCH}` but not available")

if SELECTED_ARCH in ["mamba","mamba2"]:
	LlmArchClass = MambaLMHeadModel
elif SELECTED_ARCH == "llama":
	LlmArchClass = LlamaForCausalLM
elif SELECTED_ARCH == "retnet":
	LlmArchClass = RetNetForCausalLM
else:
	raise ValueError(f"Requesting arch `{SELECTED_ARCH}` but not available")

class Model(LlmArchClass):
	def __init__(
		self,
		d_model=1024,
		n_layers=12,
		n_heads=16,
		p_dropout=0.1,

		config = cfg.model,
	):
		self.hyper_config = config
		
		hf_attention = config.attention if config is not None else None
		gradient_checkpointing = config.gradient_checkpointing if config is not None else True
		# text_tokens + rvq levels + [audio tokens * codebooks] (prom) + [audio tokens * codebooks] (resp) + stop
		vocab_size = 256 + cfg.model.max_levels + (1024 * cfg.model.max_levels) + (1024 * cfg.model.max_levels) + 1

		if SELECTED_ARCH == "llama":
			super().__init__(config=LlamaConfig(
				vocab_size=vocab_size,
				hidden_size=d_model,
				max_position_embeddings=cfg.dataset.frames_per_second * cfg.model.max_levels * 60, # max-length of 60 seconds
				intermediate_size=d_model*4,
				num_hidden_layers=n_layers,
				num_attention_heads=n_heads,
				attention_dropout=p_dropout,
				num_key_value_heads=n_heads,
				sliding_window=cfg.dataset.frames_per_second * cfg.model.max_levels * 12,
				hidden_act="gelu",
				is_encoder_decoder=False,
				is_decoder=True,
				attn_implementation=hf_attention,
			))
			
			if gradient_checkpointing:
				self.gradient_checkpointing_enable(gradient_checkpointing_kwargs=dict(
					use_reentrant=False
				))
		elif SELECTED_ARCH == "retnet":
			super().__init__(config=RetNetConfig(
				vocab_size=vocab_size,
				decoder_embed_dim=d_model,
				decoder_value_embed_dim =d_model * 2,
				decoder_retention_heads=n_heads,
				decoder_ffn_embed_dim=d_model * 4,
				decoder_layers=n_layers,
				dropout=p_dropout,
				checkpoint_activations=gradient_checkpointing,
				activation_fn="gelu",
				use_layernorm=False,
				use_biases=False,
				use_glu=True,

				#chunkwise_recurrent=self.causal and self.recurrent_chunk_size > 0,
				#recurrent_chunkwise_size=self.recurrent_chunk_size if self.causal else 0,
				#no_output_layer=True,
				#rotary_embedding_base=self.rotary_embedding_base, # 10000

				decoder_normalize_before=True,
			))
		elif SELECTED_ARCH in ["mamba","mamba2"]:
			super().__init__(config=MambaConfig(
				vocab_size=vocab_size,
				d_model=d_model,
				n_layer=n_layers*2,
				ssm_cfg={"layer": "Mamba2", "chunk_size":64} if SELECTED_ARCH == "mamba2" else {},
			))

			self.backbone.gradient_checkpointing = gradient_checkpointing

		self.accuracy_metric = MulticlassAccuracy(
			vocab_size,
			top_k=10,
			average="micro",
			multidim_average="global",
			ignore_index=-100,
		)

	def generate(
		self,
		*args,
		**kwargs
	):
		if SELECTED_ARCH in ["mamba","mamba2"]:
			kwargs["cg"] = True

			if "attention_mask" in kwargs:
				kwargs.pop("attention_mask")

			if "do_sample" in kwargs:
				kwargs.pop("do_sample")

			if "min_length" in kwargs:
				kwargs.pop("min_length")

		return super().generate(*args, **kwargs)

	def forward(
		self,
		*args,
		**kwargs,
	):
		if SELECTED_ARCH in ["mamba","mamba2"]:
			if "attention_mask" in kwargs:
				kwargs.pop("attention_mask")

		labels = kwargs.pop("labels") if "labels" in kwargs else None

		output = super().forward(*args, **kwargs)
		logits = output.logits

		# i HATE the correct way
		if labels is not None:
			if self.hyper_config is None or not self.hyper_config.loss_factors:
				loss = sum([ F.cross_entropy( logit[:-1, :], label[1:], ignore_index=-100 ) for logit, label in zip( logits, labels ) ])
				self.loss = dict(
					nll = loss,
				)

				self.stats = dict(
					acc = (sum([ self.accuracy_metric( logit, target ) for logit, target in zip( logits, labels ) ] ) / len( logits )).item()
				)

			else:
				sep = 3
				# determine specific sections to focus on
				indices = [ [ idx for idx, token in enumerate( batch ) if token == sep ] for i, batch in enumerate( labels ) ]

				text_index = 0
				resp_index = 1 # 1 indluces everything non text, -3 includes pre_resp + resp (ignores prom, probably better to include prom here)

				labels_text = [ batch[:indices[i][text_index] + 1 ] for i, batch in enumerate( labels ) ]
				labels_resp = [ batch[indices[i][resp_index] + 1:] for i, batch in enumerate( labels ) ]

				logits_text = [ batch[:indices[i][text_index] + 1 ] for i, batch in enumerate( logits ) ]
				logits_resp = [ batch[indices[i][resp_index] + 1:] for i, batch in enumerate( logits ) ]

				loss_text = sum([ F.cross_entropy( logit[:-1, :], label[1:], ignore_index=-100 ) for logit, label in zip( logits_text, labels_text ) ]) / len(logits_text) * self.hyper_config.loss_factor("text")
				loss_resp = sum([ F.cross_entropy( logit[:-1, :], label[1:], ignore_index=-100 ) for logit, label in zip( logits_resp, labels_resp ) ]) / len(logits_resp) * self.hyper_config.loss_factor("resp")

				self.loss = dict(
					text = loss_text,
					resp = loss_resp,
				)

				self.stats = dict(
					acc = dict(
						text = (sum([ self.accuracy_metric( logit, target ) for logit, target in zip( logits_text, labels_text ) ] ) / len( logits_text )).item(),
						resp = (sum([ self.accuracy_metric( logit, target ) for logit, target in zip( logits_resp, labels_resp ) ] ) / len( logits_resp )).item(),
					)
				)

		return output

def example_usage():
	cfg.trainer.backend = "local"
	cfg.hyperparameters.gradient_accumulation_steps = 1
	if cfg.audio_backend == "dac":
		cfg.sample_rate = 44_000

	from functools import partial
	from einops import repeat
	from tqdm import tqdm

	from ..emb.qnt import decode_to_file, unload_model
	from ..engines import Engine
	from ..utils import wrapper as ml
	
	import numpy as np
	import re

	device = "cuda"

	def tokenize(content):
		return torch.tensor( cfg.tokenizer.encode(content) )

	def _load_quants(path) -> Tensor:
		qnt = np.load(path, allow_pickle=True)[()]
		return torch.from_numpy(qnt["codes"].astype(np.int16))[0, :cfg.model.max_levels, :].t().to(torch.int16)

	qnt = _load_quants(f"./data/qnt.{'dac' if cfg.audio_backend == 'dac' else 'enc'}")


	text_list = [
		tokenize("ˈaɪ wɪl nˌɑːt ˈæsk ɐ sˈɛkənd tˈaɪm").to(device),
		#tokenize("ˈaɪ wɪl nˌɑːt ˈæsk ɐ sˈɛkənd tˈaɪm").to(device),
	]
	prom_list = [
		qnt[:cfg.dataset.frames_per_second, :].to(device),
		#qnt[:cfg.dataset.frames_per_second, :].to(device),
	]
	resp_list = [
		qnt[:, :].to(device),
		#qnt[cfg.dataset.frames_per_second:, :].to(device),
		#qnt[:cfg.dataset.frames_per_second, :].to(device),
	]

	text_list = text_list[:1]
	prom_list = prom_list[:1]
	resp_list = resp_list[:1]

	kwargs = {}
	model = Model(**kwargs).to(device)
	steps = 100
	if cfg.model.arch_type == "mamba2":
		steps = 100
	elif cfg.model.arch_type == "llama":
		steps = 500
	elif cfg.model.interleave:
		steps = 250

	optimizer = cfg.hyperparameters.optimizer.lower() if cfg.cfg_path is not None else "prodigy"
	scheduler = cfg.hyperparameters.scheduler.lower() if cfg.cfg_path is not None else ""
	learning_rate = cfg.hyperparameters.learning_rate if cfg.cfg_path is not None else None

	if cfg.optimizations.dadaptation:
		# do not combine the two
		if scheduler == "schedulefree":
			scheduler = ""

		learning_rate = 1.0
	
	if optimizer == "prodigy":
		if learning_rate is None:
			learning_rate = 1.0

		optimizer = ml.Prodigy
	elif optimizer == "adagrad":
		if learning_rate is None:
			learning_rate = 1.0e-2

		optimizer = ml.Adagrad
	elif optimizer == "adamw":
		if learning_rate is None:
			learning_rate = 1.0e-4

		optimizer = ml.AdamW
	elif optimizer == "sdg":
		if learning_rate is None:
			learning_rate = 1.0e-4

		optimizer = ml.SGD
	else:
		raise ValueError(f"Unrecognized optimizer: {optimizer}")

	print("Optimizer:", optimizer, "\tLearning rate:", learning_rate)

	optimizer = optimizer(model.parameters(), lr=learning_rate)

	if scheduler == "schedulefree":
		if isinstance(optimizer, ml.AdamW):
			scheduler = ml.schedulefree.AdamWScheduleFree
		elif isinstance(optimizer, ml.SGD):
			scheduler = ml.schedulefree.SGDScheduleFree
		else:
			scheduler = None

		if scheduler is not None:
			print("Scheduler:", scheduler)
			optimizer = scheduler( model.parameters(), lr = learning_rate )

	if cfg.optimizations.replace and cfg.optimizations.linear:
		model = ml.replace_linear( model )
		
	if cfg.optimizations.replace and cfg.optimizations.embedding:
		model = ml.replace_embedding( model )
	
	engine = Engine(model=model, optimizer=optimizer)

	torch.save( {
		'module': model.state_dict()
	}, f"./data/{SELECTED_ARCH}.pth" )

	print(f"{LlmArchClass} parameter count: {sum(p.numel() for p in model.parameters() if p.requires_grad)}")

	@torch.inference_mode()
	def sample( name, steps=cfg.model.max_levels*cfg.dataset.frames_per_second*6 ):
		engine.eval()
		batch_size = len(text_list)
		resp_list = None
		if cfg.model.interleave:
			input_ids, attention_mask = fold_inputs(text_list=text_list, prom_list=prom_list)
			output = model.generate(input_ids=input_ids, attention_mask=attention_mask, max_length=steps, eos_token_id=3, do_sample=False)
			
			unfolded = unfold_outputs( output )
			resp_list = unfolded["resp_list"]
		else:
			resp_list = [ [] for _ in range(batch_size) ]
			for l in range(cfg.model.max_levels):
				quant_levels = [ l for _ in range(batch_size) ]

				input_ids, attention_mask = fold_inputs(text_list=text_list, prom_list=prom_list, resp_list=resp_list, quant_levels=quant_levels)
				min_length = 1 
				for batch in input_ids:
					min_length = max( min_length, batch.shape[0] + 1 )

				output = model.generate(
					input_ids=input_ids,
					attention_mask=attention_mask,
					min_length=min_length,
					max_length=min_length+steps*2,
					eos_token_id=3,
					do_sample=False
				)
				
				unfolded = unfold_outputs( output, quant_levels=quant_levels )

				if l == 0:
					steps = 0

				for batch, resp in enumerate(unfolded["resp_list"]):
					length = resp.shape[-1]

					# store length
					if l == 0:
						steps = max( steps, length )
					# pad
					else:
						resp = resp[:steps]
						if length < steps:
							resp = torch.cat([ resp, torch.Tensor([ 0 for _ in range(steps-length) ]).to(resp) ])
					resp_list[batch].append( resp )

			for i, resp in enumerate( resp_list ):
				resp_list[i] = torch.stack( resp ).t()

		for i, batch in enumerate(resp_list):
			_ = decode_to_file(batch.to(device=device), f"data/{SELECTED_ARCH}.{cfg.audio_backend}.{i}.{name}.wav", device=device)

		unload_model()

	def train():
		engine.train()
		t = trange(steps)
		for i in t:
			stats = {"step": i}
			
			batch_size = len(text_list)
			quant_levels = None if cfg.model.interleave else torch.randint(0, cfg.model.max_levels, (batch_size,))
			if quant_levels is not None:
				resps_list = [ [] if l == 0 else resp for l, resp in zip(quant_levels, resp_list) ]
			else:
				resps_list = [ resp for resp in resp_list ]


			input_ids, attention_mask = fold_inputs(text_list=text_list, prom_list=prom_list, resp_list=resps_list, targ_list=resp_list, quant_levels=quant_levels)
			target_ids, target_attention_mask = fold_inputs(text_list=text_list, prom_list=prom_list, resp_list=resp_list, targ_list=resp_list, ignore_index=-100, quant_levels=quant_levels)
			
			stats |= engine.traverse(input_ids=input_ids, labels=target_ids, attention_mask=attention_mask)
			stats |= engine.gather_attribute("stats")
			stats |= {"grad_norm": engine.get_global_grad_norm()}

			tqdm.write(f"{stats}")

		torch.save( {
			'module': model.state_dict()
		}, f"./data/{SELECTED_ARCH}.pth" )

	#sample("init", 5)
	train()
	sample("final")

if __name__ == "__main__":
	example_usage()