"""
Copyright (c) 2025 Ynosound.
All rights reserved.

See LICENSE file in the project root for full license information.
"""


class VariableDomainSequenceOptimizer:
    """
    A class for solving sequence assignment problems with variable domains:

    We have positions i = 0..n-1, each with its own domain[i].
    We want to minimize:
        sum_{i=0}^{n-1} unary_cost(i, x_i)
      + sum_{i=0}^{n-2} binary_cost(i, x_i, i+1, x_{i+1})

    using dynamic programming, supporting different domain sizes per position.
    """

    def __init__(self, domains, unary_cost, binary_cost, precompute_binary=False):
        """
        Parameters
        ----------
        domains : list of lists
            domains[i] is the list of allowable labels for position i.
            E.g., domains[0] = [0,1,2], domains[1] = ['A','B'], etc.
        unary_cost : function (i, x) -> float
            A function that gives the cost of assigning value x at position i.
        binary_cost : function (i, x, i+1, y) -> float
            A function that gives the cost of assigning x at position i and y
            at position i+1.
        precompute_binary : bool
            If true, materialize all binary-cost tables during initialization.
            The default computes binary costs lazily during fit, which is faster
            for the small variable domains typical of MusES chord analysis.
        """
        self.n = len(domains)
        self.domains = [tuple(domain) for domain in domains]
        for i, domain in enumerate(self.domains):
            if not domain:
                raise ValueError(f"Domain at position {i} is empty")
        self.unary_cost_func = unary_cost
        self.binary_cost_func = binary_cost

        self.U = self._compute_unary_tables()

        if precompute_binary and self.n > 1:
            self.B = self._compute_binary_tables()
        elif precompute_binary:
            self.B = []
        else:
            self.B = None

        self.dp = [None] * self.n
        self.backpointer = [None] * self.n

    def _compute_unary_tables(self):
        """
        For each position i, create a list where
        U[i][d] = unary_cost_func(i, domains[i][d]).
        """
        U = []
        unary_cost = self.unary_cost_func
        for i, domain in enumerate(self.domains):
            U.append([float(unary_cost(i, label)) for label in domain])
        return U

    def _compute_binary_tables(self):
        """
        For each i in [0..n-2], create a table where
        B[i][d1][d2] = binary_cost_func(i, domains[i][d1], i+1,
        domains[i+1][d2]).
        """
        B = []
        binary_cost = self.binary_cost_func
        for i in range(self.n - 1):
            dom_i = self.domains[i]
            dom_next = self.domains[i + 1]
            B.append([
                [float(binary_cost(i, label1, i + 1, label2)) for label2 in dom_next]
                for label1 in dom_i
            ])
        return B

    def fit(self):
        """
        Run the dynamic programming to find the minimum total cost and the best assignment.

        Returns
        -------
        (min_cost, best_sequence)
          min_cost : float
              The minimal total cost.
          best_sequence : list
              A list of length n with the optimal label for each position.
        """
        if self.n == 0:
            return 0.0, []

        self.dp[self.n - 1] = list(self.U[self.n - 1])
        self.backpointer[self.n - 1] = [-1] * len(self.domains[self.n - 1])

        for i in range(self.n - 2, -1, -1):
            next_dp = self.dp[i + 1]
            unary_table = self.U[i]
            dp_i = []
            bp_i = []

            if self.B is None:
                next_domain = self.domains[i + 1]
                binary_cost = self.binary_cost_func
                for label, unary_cost in zip(self.domains[i], unary_table):
                    best_index = 0
                    best_cost = (
                        float(binary_cost(i, label, i + 1, next_domain[0]))
                        + next_dp[0]
                    )
                    for next_index in range(1, len(next_domain)):
                        cost = (
                            float(binary_cost(i, label, i + 1, next_domain[next_index]))
                            + next_dp[next_index]
                        )
                        if cost < best_cost:
                            best_cost = cost
                            best_index = next_index
                    dp_i.append(unary_cost + best_cost)
                    bp_i.append(best_index)
            else:
                binary_table = self.B[i]
                for row, unary_cost in zip(binary_table, unary_table):
                    best_index = 0
                    best_cost = row[0] + next_dp[0]
                    for next_index in range(1, len(row)):
                        cost = row[next_index] + next_dp[next_index]
                        if cost < best_cost:
                            best_cost = cost
                            best_index = next_index
                    dp_i.append(unary_cost + best_cost)
                    bp_i.append(best_index)

            self.dp[i] = dp_i
            self.backpointer[i] = bp_i

        best_start = min(range(len(self.dp[0])), key=self.dp[0].__getitem__)
        min_cost = self.dp[0][best_start]

        best_sequence = [None] * self.n
        best_sequence[0] = self.domains[0][best_start]

        prev_index = best_start
        for i in range(0, self.n - 1):
            next_index = self.backpointer[i][prev_index]
            best_sequence[i + 1] = self.domains[i + 1][next_index]
            prev_index = next_index

        return min_cost, best_sequence


# ---------------------------------------------------------------------------
# Example usage:
if __name__ == "__main__":
    # Suppose we have 4 positions, each with a different domain of labels:
    domains = [
        [0, 1],  # position 0
        [0, 1, 2],  # position 1
        ['A', 'B'],  # position 2
        [10, 20, 30]  # position 3
    ]


    # A simple unary cost function that depends on i and x
    def unary_cost(i, x):
        # e.g., cost is i * int(x != 0) just as a silly example
        # for non-integer x, we'll treat 'A'/'B' or whatever carefully
        return 1.0 if x != 0 else 0.0


    # A simple binary cost function
    def binary_cost(i, x, j, y):
        # For demonstration, let's say cost = 1 if x == y, else 0
        return float(x == y)


    optimizer = VariableDomainSequenceOptimizer(domains, unary_cost, binary_cost)
    cost, best_seq = optimizer.fit()

    print("Minimal cost:", cost)
    print("Best sequence:", best_seq)