from math import inf
from logitbias import logitbias_simplify, escape_to_zero
from logitbias_json import logitbias_eval_all


def mask_range(logits, low, high, outside=False):
    if outside:
        logits[:low] = -inf
        logits[high + 1 :] = -inf
    else:
        logits[low : high + 1] = -inf


class LogitProcessor:
    def reset(self):
        pass

    def __call__(self, ys, logits):
        raise NotImplementedError()


class SyntaxChecker(LogitProcessor):
    def __init__(
        self,
        vocab,
        min_polyphony,
        max_polyphony,
        allow_new_track=False,
        allow_end=False,
        allow_description_anchor=False,
    ):
        super().__init__()
        self.vocab = vocab
        if min_polyphony is not None:
            if min_polyphony < 0:
                raise ValueError("min_polyphony must be >= 0")
            if min_polyphony > vocab.polyphony_max:
                raise ValueError("min_polyphony must be <= vocab.polyphony_max")
            if max_polyphony is not None and max_polyphony < min_polyphony:
                raise ValueError("max_polyphony must be >= min_polyphony")
            if min_polyphony < 2:
                min_polyphony = None
        if max_polyphony is not None:
            max_polyphony = min(max_polyphony, vocab.polyphony_max)
        else:
            max_polyphony = vocab.polyphony_max

        self.min_polyphony = min_polyphony
        self.max_polyphony = max_polyphony

        self.allow_new_track = allow_new_track
        self.allow_end = allow_end
        self.allow_description_anchor = allow_description_anchor

    def __call__(self, ys, logits):
        assert ys.shape[0] > 0

        if ys.shape[0] >= 2:
            # no key symbols after header
            mask_range(
                logits,
                self.vocab.keys_least_index,
                self.vocab.keys_greatest_index,
            )

        expect_duration = False
        allow_polyphony = False

        logits[self.vocab.pad.index] = -inf
        logits[self.vocab.begin.index] = -inf
        if not self.allow_end:
            logits[self.vocab.end.index] = -inf
        if not self.allow_new_track:
            logits[self.vocab.new_track.index] = -inf
        if not self.allow_description_anchor:
            logits[self.vocab.description_anchor.index] = -inf

        lastsym = ys[-1, 0].item()

        # polyphony can be after header, new_track, description_anchor, or rest
        was_rest = ys.shape[0] > 1 and self.vocab.is_rest(ys[-2, 0].item())
        allow_polyphony = (
            self.vocab.is_begin(lastsym)
            or self.vocab.is_new_track(lastsym)
            or self.vocab.is_key(lastsym)
            or self.vocab.is_description_anchor(lastsym)
            or was_rest
        )

        # duration always and only follows pitch or rest
        expect_duration = self.vocab.is_pitch(lastsym) or self.vocab.is_rest(lastsym)
        mask_range(
            logits,
            self.vocab.durations_least_index,
            self.vocab.durations_greatest_index,
            outside=expect_duration,
        )
        if self.vocab.is_polyphony(lastsym):
            mask_range(
                logits,
                self.vocab.pitch_least_index,
                self.vocab.pitch_greatest_index,
                outside=True,
            )

        # polyphony constraints
        next_polyphony = ys[-1, 8].item()
        num_over_limit = next_polyphony - self.max_polyphony
        num_over_commitment = -1

        if allow_polyphony:
            # mask polyphony symbols that would exceed max_polyphony given current polyphony
            mask_range(
                logits,
                self.vocab.polyphony_least_index + max(1, -num_over_limit) - 1,
                self.vocab.polyphony_greatest_index,
            )
            # enforce minimum polyphony
            if self.min_polyphony is not None:
                mask_range(
                    logits,
                    self.vocab.polyphony_least_index + self.min_polyphony - 2,
                    self.vocab.polyphony_greatest_index,
                    outside=True,
                )
            # but do not mask rests!
            logits[self.vocab.rest.index] = 0
            # TODO fail fast here if infeasible
        else:
            mask_range(
                logits,
                self.vocab.polyphony_least_index,
                self.vocab.polyphony_greatest_index,
            )

            # if already committed to a polyphony, figure out how many more notes are allowed
            committed_polyphony = 1
            chord_size = 0

            for i in range(ys.shape[0] - 1, -1, -1):
                sym = ys[i, 0].item()
                if self.vocab.is_polyphony(sym):
                    committed_polyphony = self.vocab.symbols[sym].value
                    break
                if self.vocab.is_pitch(sym):
                    chord_size += 1
                if self.vocab.is_rest(sym):
                    break

            num_over_commitment = chord_size - committed_polyphony

        # only one rest in a row
        if was_rest:
            logits[self.vocab.rest.index] = -inf

        if (num_over_limit >= 0 or num_over_commitment >= 0) and not expect_duration:
            # we are at max polyphony, so prevent new pitches
            mask_range(
                logits,
                self.vocab.pitch_least_index,
                self.vocab.pitch_greatest_index,
            )

        if num_over_limit >= 0 and self.vocab.is_rest(lastsym):
            # if generating a rest duration, mask durations too short to allow generation of a pitch
            # (since the pitch would exceed max_polyphony)
            note_end_times = []
            p = next_polyphony
            t = self.vocab.get_tensor_subbeat(ys[-1], next=True).item()
            for j in range(2, ys.shape[0] + 1):
                sym = ys[-j, 0].item()
                nextsym = ys[-j + 1, 0].item()
                if self.vocab.is_pitch(sym) and self.vocab.is_duration(nextsym):
                    start_subbeat = self.vocab.get_tensor_subbeat(ys[-j]).item()
                    end_subbeat = start_subbeat + self.vocab.symbols[nextsym].value
                    if end_subbeat > t:
                        note_end_times.append(end_subbeat)
                        p -= 1
                        if p == 0:
                            break
            note_end_times.sort()
            free_up_polyphony = min(
                (
                    num_over_limit + 0
                    if self.min_polyphony is None
                    else (self.min_polyphony - 1)
                ),
                len(note_end_times) - 1,
            )
            dt = note_end_times[free_up_polyphony] - t  # brain teaser
            mask_range(
                logits,
                self.vocab.durations_least_index,
                (
                    self.vocab.durations[dt].index - 1
                    if dt <= self.vocab.duration_max
                    else self.vocab.durations_greatest_index
                ),  # this fails faster if we're already infeasible
            )


class BiasFormer(LogitProcessor):
    def __init__(self, vocab, bias, count=None):
        super().__init__()
        self.vocab = vocab
        self.bias = bias
        self.count = count
        self.env = {}
        self.last_seq_len = 0

        self.pitch_idx_vals = [
            (
                (min(v.index for v in vs.values()), max(v.index for v in vs.values())),
                p,
            )
            for p, vs in vocab.pitches.items()
        ]
        self.polyphony_idx_vals = [(p.index, p.value) for p in vocab.polyphony.values()]
        self.duration_idx_vals = [
            (p.index, p.value / vocab.quantize_divisions)
            for p in vocab.durations.values()
        ]

    def reset(self):
        self.env = {}
        self.last_seq_len = 0
        self.count = None

    def update_env(self, ys):
        if ys.shape[0] < self.last_seq_len:
            self.env = {}
            self.last_seq_len = 0
        elif ys.shape[0] == self.last_seq_len:
            return

        env = self.env
        pitches = env.get("pitches", [])
        lengths = env.get("lengths", [])
        start = env.get("start", True)
        subbeat = env.get("subbeat", -1)

        for i in range(self.last_seq_len, ys.shape[0]):
            lastval = ys[i - 1, 0].item() if i > 0 else -1
            val = ys[i, 0].item()
            val_subbeat = self.vocab.get_tensor_subbeat(ys[i]).item()
            if self.vocab.is_begin(val) or self.vocab.is_new_track(val):
                start = True
                subbeat = -1
                continue
            if not (self.vocab.is_key(val) or self.vocab.is_description_anchor(val)):
                start = False
            if (
                subbeat < val_subbeat
                and self.vocab.is_duration(val)
                and self.vocab.is_pitch(lastval)
            ):
                subbeat = val_subbeat
                pitches.append(self.vocab.symbols[lastval].value)
                lengths.append(
                    self.vocab.symbols[val].value / self.vocab.quantize_divisions
                )

        self.last_seq_len = ys.shape[0]
        self.env = {
            "start": start,
            "pitches": pitches,
            "lengths": lengths,
            "subbeat": subbeat,  # for own use
            "index": self.count if self.count is not None else 0,
            "beat": self.vocab.get_tensor_subbeat(ys[-1], next=True).item()
            / self.vocab.quantize_divisions,
        }

    def apply_parametric_bias_to_range(
        self, logits, env, variable, range_values, range_start, range_end
    ):
        if (logits[range_start : range_end + 1] == -inf).all():
            return
        parametric_bias = logitbias_simplify(self.bias, env, variable)
        if callable(parametric_bias):
            for idx, val in range_values:
                delta = parametric_bias(val)
                if isinstance(idx, tuple):
                    logits[idx[0] : idx[1] + 1] += delta
                else:
                    logits[idx] += delta
        elif parametric_bias != 0:
            logits[range_start : range_end + 1,] += parametric_bias

    def __call__(self, ys, logits):
        self.update_env(ys)
        env = self.env
        lastsym = ys[-1, 0].item() if ys.shape[1] > 0 else -1
        was_rest = self.vocab.is_rest(lastsym)
        env["length"] = escape_to_zero
        env["polyphony"] = escape_to_zero
        env["pitch"] = escape_to_zero
        env["rest"] = False

        self.apply_parametric_bias_to_range(
            logits,
            env,
            "pitch",
            self.pitch_idx_vals,
            self.vocab.pitch_least_index,
            self.vocab.pitch_greatest_index,
        )
        self.apply_parametric_bias_to_range(
            logits,
            env,
            "polyphony",
            self.polyphony_idx_vals,
            self.vocab.polyphony_least_index,
            self.vocab.polyphony_greatest_index,
        )
        env["rest"] = was_rest
        self.apply_parametric_bias_to_range(
            logits,
            env,
            "length",
            self.duration_idx_vals,
            self.vocab.durations_least_index,
            self.vocab.durations_greatest_index,
        )
        env["rest"] = True
        logits[self.vocab.rest.index] += logitbias_eval_all(self.bias, env)


class SearchException(Exception):
    pass