import numpy as np
import torch
from math import inf
from logitbias_json import json2logitbias
from sample import (
    BiasFormer,
    SearchException,
    SyntaxChecker,
)
from dataclasses import dataclass
from typing import Optional, List, Any
from logitbias import eval as logitbias_eval, VariableStub
from logitbias_parser import parse as logitbias_parse
from threading import Thread
from queue import Queue, Empty
from contextlib import contextmanager
from loader import load_by_path
from lib.perf_recorder import BasePerfRecorder, PerfRecorder
import time

# TODO placement of description anchor token for inference!


class RedundancyChecker:
    def __init__(
        self,
        vocab,
        pitch_penalty: float,
        duration_penalty: float,
        own_symbols: np.ndarray,
        clip_refs: np.ndarray,
        ctx_refs: np.ndarray,
    ):
        self.vocab = vocab
        self.pitch_penalty = pitch_penalty
        self.duration_penalty = duration_penalty
        self.own_symbols = own_symbols
        self.clip_refs = clip_refs
        self.ctx_refs = ctx_refs
        self.reset()

    def before_sample(self, pos, logits):
        if self.all_distinct:
            return

        # evaluate distinctness over [self.pos, pos) already-sampled symbols
        for i in range(self.pos, pos):
            sym = self.own_symbols[i].item()
            assert sym != self.vocab.pad.index

            if self.first_pitch_pos == -1 and self.vocab.is_pitch(sym):
                self.first_pitch_pos = i
            if self.first_pitch_pos != -1:
                assert self.first_pitch_pos <= i
                nd_syms = self.clip_refs[self.clip_refs_nd, i - self.first_pitch_pos]
                nd_pred = nd_syms == sym
                self.clip_refs_nd = self.clip_refs_nd[nd_pred]

            nd_syms = self.ctx_refs[self.ctx_refs_nd, i]
            nd_pred = nd_syms == sym
            self.ctx_refs_nd = self.ctx_refs_nd[nd_pred]
        self.pos = pos

        if self.ctx_refs_nd.shape[0] == 0 and self.clip_refs_nd.shape[0] == 0:
            self.all_distinct = True
            return

        # penalize logits matching next symbol in non-distinct clips + ctxs
        nd_syms = np.unique(
            np.concatenate(
                [
                    self.clip_refs[
                        self.clip_refs_nd,
                        (
                            (pos - self.first_pitch_pos)
                            if self.first_pitch_pos != -1
                            else 0
                        ),
                    ].reshape(-1),
                    self.ctx_refs[self.ctx_refs_nd, pos].reshape(-1),
                ]
            )
        )
        pitch_syms = nd_syms[
            np.logical_or(
                np.logical_and(
                    nd_syms >= self.vocab.pitch_least_index,
                    nd_syms < self.vocab.pitch_greatest_index,
                ),
                np.logical_and(
                    nd_syms >= self.vocab.polyphony_least_index,
                    nd_syms < self.vocab.polyphony_greatest_index,
                ),
            )
        ]
        duration_syms = nd_syms[
            np.logical_and(
                nd_syms >= self.vocab.durations_least_index,
                nd_syms < self.vocab.durations_greatest_index,
            )
        ]

        logits[pitch_syms] -= self.pitch_penalty
        logits[duration_syms] -= self.duration_penalty

    def reset(self):
        self.first_pitch_pos = -1
        self.pos = 0
        self.all_distinct = False
        self.clip_refs_nd = np.arange(0, self.clip_refs.shape[0], dtype=np.int32)
        self.ctx_refs_nd = np.arange(0, self.ctx_refs.shape[0], dtype=np.int32)


def maybe_coerce(typ, x):
    if x is None:
        return None
    return typ(x)


@dataclass
class GenerationRequest:
    temperature: float
    stop_cond: str
    bias: Any = None
    min_polyphony: Optional[int] = None
    max_polyphony: Optional[int] = None
    dists: bool = False
    rest_temperature: float = None
    note_length_temperature: float = None
    cfg_scale: float = 1.5

    def __post_init__(self):
        if self.rest_temperature is None:
            self.rest_temperature = self.temperature
        if self.note_length_temperature is None:
            self.note_length_temperature = self.temperature

    @classmethod
    def from_json(cls, obj):
        bias = obj.get("bias", None)
        if bias is not None:
            bias = json2logitbias(bias)
        if "stopCondition" in obj:
            if "count" in obj:
                raise RuntimeError("Cannot specify both stopCondition and count")
            stop_cond = str(obj["stopCondition"])
        else:
            stop_cond = f"steps == {int(obj.get('count', 5))}"

        # for reverse compatibility
        timing_temperature = obj.get("timingTemperature", None)

        return cls(
            temperature=float(obj.get("temperature", 0.9)),
            stop_cond=stop_cond,
            bias=bias,
            min_polyphony=maybe_coerce(int, obj.get("minPolyphony", None)),
            max_polyphony=maybe_coerce(int, obj.get("maxPolyphony", None)),
            dists=maybe_coerce(bool, obj.get("dists", None)),
            rest_temperature=maybe_coerce(
                float, obj.get("restTemperature", timing_temperature)
            ),
            note_length_temperature=maybe_coerce(
                float, obj.get("noteLengthTemperature", timing_temperature)
            ),
            cfg_scale=float(obj.get("cfgScale", 1.5)),
        )


@dataclass
class GenerationRequestCtx:
    stop_cond: Any
    bias_former: Optional[BiasFormer]
    logit_processors: List[Any]
    request: GenerationRequest
    idx: int
    cfg_idx: int
    redundancy_checker: Optional[RedundancyChecker]

    stop_pos: int = -1
    step_count_pos: int = 0
    in_chord: bool = False
    step_count: int = 0

    backtrack_count: int = 0
    backtrack_from_pos: int = -1
    pending_search_error: bool = False
    backtrack_total_count: int = 0


@dataclass
class SearchCtx:
    request_ctxts: List[GenerationRequestCtx]
    encoder_input_ids: Optional[torch.Tensor]
    encoder_attention_mask: Optional[torch.Tensor]
    batch_size: int
    start: torch.Tensor
    logit_masks: np.ndarray
    ys: np.ndarray
    dists: np.ndarray
    pos: int
    recorder: BasePerfRecorder
    start_time: float
    deadline: float

    logit_processor_time: float = 0.0
    sampling_time: float = 0.0
    redundancy_checker_time: float = 0.0
    stop_cond_time: float = 0.0
    logits0: Optional[torch.Tensor] = None
    logits: Any = None


@contextmanager
def perf_stopwatch(recorder, name):
    start = time.time()
    yield
    recorder.record(name, 1000.0 * (time.time() - start), "ms")


class ComposerAPI:
    @classmethod
    def from_path(cls, path):
        vocab, model, critic = load_by_path(path)
        return cls(vocab, model, critic)

    def __init__(self, vocab, model, critic=None, backtrack_limit=10):
        super().__init__()
        self.vocab = vocab
        self.model = model
        self.critic = critic
        self.backtrack_limit = backtrack_limit
        self.max_seq_len = model.params.max_inference_seq_len
        # self.prompt_prefix = "tags: [analyzed aligned] "  # TODO cleanup - also this must match data_aug.py
        self.prompt_prefix = ""
        self.inference_batch_size = model.params.inference_batch_size

    # TODO accompany becomes list of midi
    def generate(self, *args, recorder=PerfRecorder(), **kwargs):
        with torch.no_grad(), perf_stopwatch(recorder, "request"):
            return self.generate_impl(*args, recorder=recorder, **kwargs)

    def setup_request_context(
        self,
        idx: int,
        count: int,
        request: GenerationRequest,
        redundant_pitch_penalty: float,
        redundant_duration_penalty: float,
        redundant_clips: np.ndarray,
        ys: np.ndarray,
        start_len: int,
    ):
        stop_cond = logitbias_parse(request.stop_cond)
        stop_cond = logitbias_eval(
            stop_cond,
            {k: VariableStub() for k in ["steps", "error", "maxProbability", "beat"]},
        )

        logit_processors = []

        logit_processors.append(
            SyntaxChecker(
                self.vocab,
                request.min_polyphony,
                request.max_polyphony,
            )
        )

        bias_former = None
        if request.bias is not None:
            bias_former = BiasFormer(self.vocab, request.bias)
            logit_processors.append(bias_former)

        if redundant_pitch_penalty > 0.0 or redundant_duration_penalty > 0.0:
            redundancy_checker = RedundancyChecker(
                self.vocab,
                redundant_pitch_penalty,
                redundant_duration_penalty,
                ys[idx, start_len:, 0],
                redundant_clips,
                ys[:idx, start_len:, 0],
            )
        else:
            redundancy_checker = None

        return GenerationRequestCtx(
            stop_cond,
            bias_former,
            logit_processors,
            request,
            idx,
            count + idx,
            redundancy_checker,
        )

    @contextmanager
    def forward_pass_thread(self, recorder: BasePerfRecorder):
        threaded_forward_q_in: Queue = Queue(1)
        threaded_forward_q_out: Queue = Queue()
        logits_bottleneck_count = 0
        model_bottleneck_count = 0
        model_forward_time = 0.0

        def threaded_forward_entry():
            nonlocal model_forward_time
            while True:
                req = threaded_forward_q_in.get()
                if req is None:
                    break
                try:
                    start = time.time()
                    res = self.model.forward(**req)
                    if callable(res):
                        res = res()
                    res = res.cpu().detach()
                    model_forward_time += time.time() - start
                    threaded_forward_q_out.put(res)
                except Exception as exn:  # pylint: disable=broad-except
                    threaded_forward_q_out.put(exn)

        threaded_forward_thd = Thread(target=threaded_forward_entry, daemon=True)
        threaded_forward_thd.start()

        def threaded_forward(**kwargs):
            threaded_forward_q_in.put(kwargs)

            def thunk():
                nonlocal logits_bottleneck_count
                nonlocal model_bottleneck_count
                try:
                    res = threaded_forward_q_out.get(block=False)
                    logits_bottleneck_count += 1
                except Empty:
                    res = threaded_forward_q_out.get()
                    model_bottleneck_count += 1
                if isinstance(res, Exception):
                    raise res
                return res

            return thunk

        try:
            yield threaded_forward
        finally:
            threaded_forward_q_in.put(None)
            threaded_forward_thd.join()
        recorder.record("threaded_forward", 1000.0 * model_forward_time, "msThd")
        recorder.record("logits_bottleneck", logits_bottleneck_count, "count")
        recorder.record("model_bottleneck", model_bottleneck_count, "count")

    def search_step(self, threaded_forward, sctx: SearchCtx):
        logit_processor_start = time.time()

        running_ctxts = [c for c in sctx.request_ctxts if c.stop_pos == -1]

        if len(running_ctxts) == 0:
            return False

        if sctx.pos >= sctx.ys.shape[1]:
            for c in running_ctxts:
                sctx.recorder.log(
                    f"search reached max length {sctx.pos} for context {c.idx}"
                )
                c.stop_pos = sctx.pos
            return False

        if time.time() >= sctx.deadline:
            sctx.recorder.log(f"search reached deadline at pos {sctx.pos}")
            for c in running_ctxts:
                c.stop_pos = sctx.pos
            return False

        # Evaluate logit processors for non-done contexts to get masks
        for ctx in running_ctxts:
            sctx.logit_masks[ctx.idx].fill(0.0)
            for logit_processor in ctx.logit_processors:
                logit_processor(sctx.ys[ctx.idx, : sctx.pos], sctx.logit_masks[ctx.idx])

            if ctx.backtrack_from_pos != -1 and sctx.pos > ctx.backtrack_from_pos:
                sctx.recorder.log(
                    f"backtrack: progress for context {ctx.idx} after {ctx.backtrack_count} attempts"
                )
                ctx.backtrack_from_pos = -1

        sctx.logit_processor_time += time.time() - logit_processor_start

        if callable(sctx.logits):
            sctx.logits = sctx.logits()[:, -1]

        sampling_start = time.time()

        # CFG
        if self.model.encoder is not None:
            for ctx in running_ctxts:
                logits_uncond = sctx.logits[ctx.cfg_idx]
                sctx.logits[ctx.idx] = logits_uncond + ctx.request.cfg_scale * (
                    sctx.logits[ctx.idx] - logits_uncond
                )

        # TODO this only needs to be done for running contexts
        sctx.logits += torch.from_numpy(sctx.logit_masks)
        probs = torch.softmax(sctx.logits, dim=-1).numpy()
        sctx.dists[:, sctx.pos] = probs

        # Sample next token for non-done contexts
        backtracking = False
        backtracking_to_start = False
        for ctx in running_ctxts:
            if torch.all(sctx.logits[ctx.idx] < -1000.0):
                # backtrack
                if ctx.backtrack_from_pos == -1:
                    sctx.recorder.log(
                        f"backtrack: starting from {sctx.pos} for context {ctx.idx}"
                    )
                    ctx.backtrack_from_pos = sctx.pos
                    ctx.backtrack_count = 0
                if ctx.backtrack_count > self.backtrack_limit:
                    ctx.pending_search_error = True
                    continue  # skip sampling
                else:
                    sctx.pos = max(
                        sctx.start.shape[0],
                        ctx.backtrack_from_pos - 4 * int(pow(1.5, ctx.backtrack_count)),
                    )
                    backtracking = True
                    if sctx.pos == sctx.start.shape[0]:
                        backtracking_to_start = True
                    ctx.backtrack_total_count += 1
                    ctx.backtrack_count += 1
                    for ctx2 in running_ctxts:
                        ctx2.step_count_pos = 0
                        if ctx2.bias_former is not None:
                            ctx2.bias_former.count = 0
                        if ctx2.redundancy_checker is not None:
                            ctx2.redundancy_checker.reset()
                        sctx.ys[ctx2.idx, sctx.pos + 1 :, 0] = 0
                    break

            if ctx.redundancy_checker is not None:
                start = time.time()
                ctx.redundancy_checker.before_sample(
                    sctx.pos - sctx.start.shape[0], sctx.logits[ctx.idx]
                )
                sctx.redundancy_checker_time += time.time() - start

            temperature = ctx.request.temperature
            if self.vocab.is_rest(sctx.ys[ctx.idx, sctx.pos - 1, 0]):
                temperature = ctx.request.rest_temperature
            if self.vocab.is_pitch(sctx.ys[ctx.idx, sctx.pos - 1, 0]):
                temperature = ctx.request.note_length_temperature
            temperature = float(temperature)

            if temperature == 0.0:
                next_token = torch.argmax(sctx.logits[ctx.idx])
            else:
                next_token = torch.multinomial(
                    torch.softmax(sctx.logits[ctx.idx] / temperature, dim=-1),
                    num_samples=1,
                )

            self.vocab.append_symbols(
                sctx.ys[ctx.idx : ctx.idx + 1],
                next_token.view(1, 1),
                strict=False,
                seq_len=sctx.pos,
            )
            if self.model.encoder is not None:
                sctx.ys[ctx.cfg_idx, sctx.pos] = sctx.ys[ctx.idx, sctx.pos]

        sctx.sampling_time += time.time() - sampling_start

        if backtracking_to_start:
            sctx.logits = sctx.logits0
            return True

        # Kick off model forward pass -> logits
        sctx.logits = threaded_forward(
            x=torch.from_numpy(sctx.ys[:, sctx.pos : sctx.pos + 1]),
            start_pos=sctx.pos,
            encoder_input_ids=None,
            encoder_attention_mask=None,
        )

        stop_cond_start = time.time()

        # Evaluate stop conditions for non-done contexts
        if not backtracking:
            for ctx in running_ctxts:
                if ctx.step_count_pos == 0:
                    ctx.step_count = 0
                    ctx.in_chord = False
                for j in range(sctx.start.shape[0] + ctx.step_count_pos, sctx.pos + 1):
                    sym = sctx.ys[ctx.idx, j, 0].item()
                    if self.vocab.is_pitch(sym) and not ctx.in_chord:
                        ctx.in_chord = True
                    if self.vocab.is_rest(sym) and ctx.in_chord:
                        ctx.in_chord = False
                        ctx.step_count += 1
                ctx.step_count_pos = max(sctx.pos + 1 - sctx.start.shape[0], 0)
                if ctx.bias_former is not None:
                    ctx.bias_former.count = ctx.step_count

                if self.vocab.embedding_time_wrapped(sctx.ys[ctx.idx, sctx.pos]):
                    sctx.recorder.log(f"context {ctx.idx} wrapped around")
                    ctx.pending_search_error = True

                max_probability = np.max(probs[ctx.idx])
                env = {
                    "steps": ctx.step_count,
                    "error": ctx.pending_search_error,
                    "maxProbability": max_probability.item(),
                    "beat": self.vocab.get_tensor_subbeat(
                        sctx.ys[ctx.idx, sctx.pos - 1], next=True
                    ).item()
                    / self.vocab.quantize_divisions,
                    "generationTime": time.time() - sctx.start_time,
                }
                if logitbias_eval(ctx.stop_cond, env):
                    sctx.recorder.log(
                        f"stop: context {ctx.idx} at pos {sctx.pos} with env {env}"
                    )
                    ctx.stop_pos = sctx.pos
                elif ctx.pending_search_error:
                    raise SearchException(
                        f"search failed for context {ctx.idx} at pos {sctx.pos}"
                    )

        sctx.stop_cond_time += time.time() - stop_cond_start
        sctx.pos += 1
        return True

    def generate_ctx(self, sctx: SearchCtx):
        with self.forward_pass_thread(sctx.recorder) as threaded_forward:
            # common forward pass
            with perf_stopwatch(sctx.recorder, "common_forward"):
                rearranger = [0] * len(sctx.request_ctxts)
                if self.model.encoder is not None:
                    rearranger += [1] * len(sctx.request_ctxts)
                rearranger = torch.tensor(rearranger, dtype=torch.long)
                self.model.hint_inference_batch_size(sctx.batch_size)
                sctx.logits0 = (
                    self.model.forward(
                        x=torch.from_numpy(sctx.ys[0, : sctx.start.shape[0]])
                        .unsqueeze(0)
                        .repeat(2 if self.model.encoder is not None else 1, 1, 1),
                        start_pos=0,
                        encoder_input_ids=sctx.encoder_input_ids,
                        encoder_attention_mask=sctx.encoder_attention_mask,
                    )[:, -1]
                    .cpu()
                    .detach()[rearranger]
                )
                self.model.rearrange(rearranger)
                sctx.logits = sctx.logits0

            with perf_stopwatch(sctx.recorder, "search"):
                while self.search_step(threaded_forward, sctx):
                    pass

    def generate_impl(
        self,
        prefix,
        generation_requests,
        recorder,
        accompany=None,
        text_prompt=None,
        deadline=inf,
        redundant_pitch_penalty=0.0,
        redundant_duration_penalty=0.0,
        redundant_clips=[],
    ):
        start_time = time.time()

        assert redundant_pitch_penalty >= 0.0 and redundant_duration_penalty >= 0.0

        if self.model.encoder is None:
            assert text_prompt is None
            encoder_input_ids, encoder_attention_mask = None, None
        else:
            with perf_stopwatch(recorder, "tokenize"):
                encoder_input_ids, encoder_attention_mask = self.model.encoder.tokenize(
                    [
                        self.prompt_prefix
                        + (text_prompt if text_prompt is not None else ""),
                        "",  # for CFG
                    ]
                )

        with perf_stopwatch(recorder, "midi_to_tensor"):
            start = self.vocab.midi_to_tensor(
                prefix, accompany=accompany, rest_end=True
            )

        with perf_stopwatch(recorder, "redundant_clips_to_tensor"):
            redundant_clip_tensor = np.zeros(
                (len(redundant_clips), self.max_seq_len), dtype=np.int32
            )
            for i, clip in enumerate(redundant_clips):
                clip_tensor = self.vocab.midi_to_tensor(clip)
                for j in range(clip_tensor.shape[0]):
                    if self.vocab.is_pitch(clip_tensor[j, 0].item()):
                        insert_len = min(self.max_seq_len, clip_tensor.shape[0] - j)
                        redundant_clip_tensor[i, :insert_len] = clip_tensor[
                            j : insert_len + j, 0
                        ]
                        break

        recorder.record("start", start.shape[0], "count")

        if start.shape[0] >= self.max_seq_len - 4:
            raise RuntimeError("input too long for model")

        start_midi, _ = self.vocab.tensor_to_midi(start, strict=False)

        ctxts = []
        batch_size_per_request = 1 if self.model.encoder is None else 2

        if (
            batch_size_per_request * len(generation_requests)
            > self.inference_batch_size
        ):
            raise RuntimeError("too many generation requests")

        if self.critic is None:
            runs_per_request = 1
        else:
            runs_per_request = self.inference_batch_size // (
                batch_size_per_request * len(generation_requests)
            )

        num_ctxs = len(generation_requests) * runs_per_request
        batch_size = num_ctxs * batch_size_per_request

        with perf_stopwatch(recorder, "alloc"):
            # allocate working tensors
            ys = np.zeros(
                (batch_size, self.max_seq_len, 9),
                dtype=np.int32,
            )
            dists = np.zeros(
                (batch_size, self.max_seq_len, self.vocab.N),
                dtype=np.float32,
            )
            # faster for complex indexing
            logit_masks = np.zeros((batch_size, self.vocab.N), dtype=np.float32)
            ys[:, : start.shape[0]] = start.unsqueeze(0)

        for i, request in enumerate(generation_requests):
            with perf_stopwatch(recorder, f"setup_context[{i}]"):
                for j in range(runs_per_request):
                    ctxts.append(
                        self.setup_request_context(
                            i * runs_per_request + j,
                            num_ctxs,
                            request,
                            redundant_pitch_penalty,
                            redundant_duration_penalty,
                            redundant_clip_tensor,
                            ys,
                            start.shape[0],
                        )
                    )

        search_ctx = SearchCtx(
            batch_size=batch_size,
            request_ctxts=ctxts,
            encoder_input_ids=encoder_input_ids,
            encoder_attention_mask=encoder_attention_mask,
            start=start,
            logit_masks=logit_masks,
            ys=ys,
            dists=dists,
            pos=start.shape[0],
            recorder=recorder,
            start_time=start_time,
            deadline=deadline,
        )

        self.generate_ctx(search_ctx)

        recorder.record(
            "logit_processors", 1000.0 * search_ctx.logit_processor_time, "msThd"
        )
        recorder.record("stop_cond", 1000.0 * search_ctx.stop_cond_time, "msThd")
        recorder.record("sampling", 1000.0 * search_ctx.sampling_time, "ms")
        recorder.record(
            "redundancy_checker", 1000.0 * search_ctx.redundancy_checker_time, "ms"
        )
        recorder.record("steps", search_ctx.pos - start.shape[0], "count")

        if self.critic is None:
            scores = torch.tensor([0.0] * len(ctxts))
        else:
            with perf_stopwatch(search_ctx.recorder, "embeddings"):
                max_stop_pos = max(ctx.stop_pos for ctx in ctxts)
                embeddings = (
                    self.model.forward(
                        x=torch.from_numpy(search_ctx.ys[:, : max_stop_pos + 1]),
                        start_pos=0,
                        encoder_input_ids=search_ctx.encoder_input_ids,
                        encoder_attention_mask=search_ctx.encoder_attention_mask,
                        return_embeddings=True,
                    )
                    .cpu()
                    .detach()
                )
                scores = self.critic(embeddings).cpu().detach().view(-1)

        with perf_stopwatch(recorder, "tensor_to_midi"):
            res = []
            for i, req in enumerate(generation_requests):
                group = [c for c in ctxts if id(c.request) == id(req)]
                recorder.record(
                    f"backtracks[{i}]",
                    sum(c.backtrack_total_count for c in group),
                    "count",
                )
                best_ctx = max(group, key=lambda c: scores[c.idx])
                midi, _ = self.vocab.tensor_to_midi(
                    ys[best_ctx.idx, : best_ctx.stop_pos + 1],
                    strict=False,
                    dists=(
                        dists[best_ctx.idx, start.shape[0] : best_ctx.stop_pos + 1]
                        if req.dists
                        else None
                    ),
                )
                res.append(midi[0 if start_midi is None else len(start_midi) :])

        return res