import numpy as np
import mlx.core as mx
import mlx.nn as nn
from mlx_lm.sample_utils import make_sampler
import mlx.optimizers as optim
from ..inference_engine import InferenceEngine
from .sharded_utils import load_model_shard, resolve_tokenizer
from .losses import loss_fns
from ..shard import Shard
from typing import Optional
from .download.shard_download import ShardDownloader
import asyncio
from collections import OrderedDict
from mlx_lm.models.cache import make_prompt_cache
from concurrent.futures import ThreadPoolExecutor


class MLXDynamicShardInferenceEngine(InferenceEngine):
    def __init__(self, shard_downloader: ShardDownloader):
        self.shard = None
        self.shard_downloader = shard_downloader
        self.caches = OrderedDict()
        self.sampler_params: tuple[float, float] = (0.0, 0.0, 0.0, 1)
        self.sampler = make_sampler(*self.sampler_params)
        self._mlx_thread = ThreadPoolExecutor(max_workers=1, thread_name_prefix="mlx")
        self._tokenizer_thread = ThreadPoolExecutor(
            max_workers=1, thread_name_prefix="tokenizer"
        )
        self.session = {}
        self._shard_lock = asyncio.Lock()

    async def _eval_mlx(self, *args):
        await asyncio.get_running_loop().run_in_executor(
            self._mlx_thread, mx.eval, *args
        )

    async def poll_state(self, request_id: str, max_caches=2):
        if request_id in self.caches:
            self.caches.move_to_end(request_id)
        else:
            newcache = make_prompt_cache(self.model)
            if len(self.caches) > max_caches:
                self.caches.popitem(last=False)
            self.caches[request_id] = newcache
        return {"cache": self.caches[request_id]}

    async def sample(
        self, x: np.ndarray, temp: float = 0.0, top_p: float = 1.0
    ) -> np.ndarray:
        if (temp, top_p, 0.0, 1) != self.sampler_params:
            self.sampler_params = (temp, top_p, 0.0, 1)
            self.sampler = make_sampler(*self.sampler_params)
        logits = mx.array(x)
        logits = logits[:, -1, :]
        logprobs = logits - mx.logsumexp(logits, keepdims=True)
        result = self.sampler(logprobs)
        await self._eval_mlx(result)
        return np.asarray(result, dtype=int)

    async def encode(self, shard: Shard, prompt: str) -> np.ndarray:
        await self.ensure_shard(shard)
        return np.asarray(
            await asyncio.get_running_loop().run_in_executor(
                self._tokenizer_thread, self.tokenizer.encode, prompt
            )
        )

    async def decode(self, shard: Shard, tokens) -> str:
        await self.ensure_shard(shard)
        return await asyncio.get_running_loop().run_in_executor(
            self._tokenizer_thread, self.tokenizer.decode, tokens
        )

    async def save_checkpoint(self, shard: Shard, path: str):
        await self.ensure_shard(shard)
        await asyncio.get_running_loop().run_in_executor(
            self._mlx_thread, lambda: self.model.save_weights(path)
        )

    async def load_checkpoint(self, shard: Shard, path: str):
        await self.ensure_shard(shard)
        await asyncio.get_running_loop().run_in_executor(
            self._mlx_thread, lambda: self.model.load_weights(path)
        )

    async def infer_tensor(
        self,
        request_id: str,
        shard: Shard,
        input_data: np.ndarray,
        inference_state: Optional[dict] = None,
    ) -> tuple[np.ndarray, Optional[dict]]:
        await self.ensure_shard(shard)
        state = (
            await self.poll_state(request_id)
            if self.model.model_type != "StableDiffusionPipeline"
            else {}
        )
        x = mx.array(input_data)

        if self.model.model_type != "StableDiffusionPipeline":
            output_data = await asyncio.get_running_loop().run_in_executor(
                self._mlx_thread,
                lambda: self.model(x, **state, **(inference_state or {})),
            )
            inference_state = None
        else:
            result = await asyncio.get_running_loop().run_in_executor(
                self._mlx_thread,
                lambda: self.model(x, **state, **(inference_state or {})),
            )
            output_data, inference_state = result

        await self._eval_mlx(output_data)
        output_data = await asyncio.get_running_loop().run_in_executor(
            self._mlx_thread, lambda: np.array(output_data, copy=False)
        )
        return output_data, inference_state

    async def evaluate(
        self,
        request_id: str,
        shard: Shard,
        inputs,
        targets,
        lengths,
        loss: str = "length_masked_ce",
    ):
        await self.ensure_shard(shard)
        await self.save_session("loss", loss_fns[loss])
        x = mx.array(inputs)
        y = mx.array(targets)
        l = mx.array(lengths)

        score = await asyncio.get_running_loop().run_in_executor(
            self._mlx_thread, lambda: self.session["loss"](self.model, x, y, l)
        )
        return score

    async def ensure_train(
        self,
        shard: Shard,
        loss: str,
        opt=optim.SGD,
        lr=1e-5,
        trainable_layers=["input_layernorm", "gate_proj"],
    ):
        await self.ensure_shard(shard)

        if (
            "train_layers" not in self.session
            or self.session["train_layers"] != trainable_layers
        ):
            await self.save_session("train_layers", trainable_layers)

            def freeze_unfreeze():
                self.model.freeze()
                self.model.apply_to_modules(
                    lambda k, v: (
                        v.unfreeze()
                        if any(
                            k.endswith(layer_name) for layer_name in trainable_layers
                        )
                        else None
                    )
                )

            await asyncio.get_running_loop().run_in_executor(
                self._mlx_thread, freeze_unfreeze
            )

        if (
            "lossname" not in self.session
            or "LVaG" not in self.session
            or self.session["lossname"] != loss
        ):
            await self.save_session("lossname", loss)
            await self.save_session(
                "LVaG", nn.value_and_grad(self.model, loss_fns[loss])
            )

        if "opt" not in self.session:
            await self.save_session("opt", opt(lr))
        return True

    async def train(
        self,
        request_id: str,
        shard: Shard,
        inputs,
        targets,
        lengths,
        loss: str = "length_masked_ce",
        opt=optim.SGD,
        lr=1e-5,
    ):
        await self.ensure_train(shard, loss, opt, lr)

        def train_step(inp, tar, lng):
            lval, grad = self.session["LVaG"](self.model, inp, tar, lng)
            gradlayers = grad["model"]["layers"]
            self.session["opt"].update(self.model, grad)
            return (
                lval,
                gradlayers,
                (self.model.parameters(), self.session["opt"].state, lval),
            )

        x = mx.array(inputs)
        y = mx.array(targets)
        l = mx.array(lengths)
        score, gradients, eval_args = await asyncio.get_running_loop().run_in_executor(
            self._mlx_thread, lambda: train_step(x, y, l)
        )
        await self._eval_mlx(*eval_args)

        layers = [
            {k: v["weight"] for k, v in layer.items() if "weight" in v}
            for layer in gradients
            if layer
        ]
        first_layer = np.array(layers[0]["input_layernorm"], copy=False)
        await self._eval_mlx(first_layer)
        return score, first_layer

    async def ensure_shard(self, shard: Shard):
        async with self._shard_lock:
            if self.shard == shard:
                return
            model_path = await self.shard_downloader.ensure_shard(
                shard, self.__class__.__name__
            )
            if self.shard != shard:
                model_shard = await asyncio.get_running_loop().run_in_executor(
                    self._mlx_thread,
                    lambda: load_model_shard(model_path, shard, lazy=False),
                )
                if hasattr(model_shard, "tokenizer"):
                    self.tokenizer = model_shard.tokenizer
                else:
                    self.tokenizer = await resolve_tokenizer(model_path)
                self.shard = shard
                self.model = model_shard
                self.caches = OrderedDict()
                self.session = {}

    async def cleanup(self):
        self._mlx_thread.shutdown(wait=True)
