"""Spherical unitarity test for classical groups: port of spherical.pm.

Given a type A, B, C or D and a parameter lambda, :func:`process`
decides whether the spherical representation X(lambda) of the split
real or p-adic group is unitary, following Dan Barbasch's algorithm:

* lambda is split into its integral (classical) part and the rest
  (:func:`sortIntegerClasses`, :func:`getClassicalInfo`);
* this gives a nilpotent orbit O on the dual side, a Levi factor
  M = M_0 x GL(k_1) x ... and a decomposition lambda = h + nu
  (:func:`moveToGL`, :func:`getGLInfo`);
* nu is a parameter for the centralizer Z of O (:func:`centralizer`),
  and X(lambda) is unitary iff nu is in the 0-complementary series of
  each factor of Z (:func:`test`, :func:`testComplementarySeries`).

Values are Perl-like (see :mod:`perlval`) so that output matches the
Perl original.
"""

import re
from functools import cmp_to_key

from .perlval import (PFrac, frac, pstr, pnum, padd, psub, pmul, pdiv,
                      pneg, pabs, pcmp, peq, plt, ple, pgt, pge, pint,
                      fix_str, n_div)
from .rootsystem import (fundamentalWeights, matrixMult, transpose,
                         simpleCoRoots)

verbose = 0


class Stop(Exception):
    pass


class Result(object):
    """What process() returns (the Perl returned a long list)."""

    def __init__(self, **kw):
        self.type = None
        self.lam = []
        self.lambdaWeight = []
        self.lambdaWeightGap = []
        self.result = 0
        self.psd = {}
        self.message = None
        self.Z = {}
        self.O = []
        self.O_0 = []
        self.M_0 = []
        self.h_0 = []
        self.O_GL = []
        self.M_GL = []
        self.h_GL = []
        self.nu_GL = []
        self.nu = {}
        self.__dict__.update(kw)


def _key(cmpfun):
    return cmp_to_key(cmpfun)


def descending(a, b):
    """Sort by the fraction value, largest first."""
    return pcmp(frac(b), frac(a))


def ascending(a, b):
    return pcmp(frac(a), frac(b))


def sort_desc(values):
    """Perl: sort descending @values (stable; fractions parsed once)."""
    dec = [(frac(v), v) for v in values]
    dec.sort(key=_key(lambda a, b: pcmp(b[0], a[0])))
    return [v for _, v in dec]


def sort_asc(values):
    """Perl: sort ascending @values."""
    dec = [(frac(v), v) for v in values]
    dec.sort(key=_key(lambda a, b: pcmp(a[0], b[0])))
    return [v for _, v in dec]


def _hkey(x):
    """A Perl hash key (the stringified value)."""
    return pstr(x)


def process(type, arg, coordinates="", out=None):
    """Test one parameter lambda.

    ``arg`` is the list of coordinates (strings such as "3/2" or ".75",
    or numbers); ``coordinates`` is 'standard' (or ''), 'bourbaki' or
    'gap'.  Text output (verbose > 0, as in the command line version)
    is appended to ``out``.
    """
    if out is None:
        out = []
    p = out.append
    lam = []
    message = None
    for i in arg:
        if isinstance(i, str) and "/" in i:
            a, b = (i.split("/") + [""])[:2]
            i = frac(a, b)
        lam.append(i)
    if not lam:
        return None
    lambdaWeight = []
    lambdaWeightGap = []
    if re.search("bourbaki|gap", coordinates or ""):
        lambdaWeight = lam
        rank = len(lambdaWeight)
        gap = "gap" if coordinates == "gap" else ""
        fws = fundamentalWeights(type, rank, gap)
        lam = matrixMult([lam], fws)[0]
        lam = dominant(type, lam)
        lambdaWeightGap = list(lambdaWeight)
        if type != "A":
            if coordinates == "bourbaki":
                lambdaWeightGap = list(reversed(lambdaWeight))
            else:
                lambdaWeightGap = list(lambdaWeight)
                lambdaWeight = list(reversed(lambdaWeightGap))
    else:
        if type == "A":
            lam = sorted(lam, key=_key(lambda a, b: pcmp(b, a)))
            if pge(lam[-1], 0) and pgt(lam[0], 0):
                message = ("All coordinates of lambda >=0, type A:\n"
                           "Assuming they should all be repeated with "
                           "their negatives")
                neg = []
                for x in lam:
                    neg.insert(0, pneg(x))
                lam = lam + neg
            else:
                last = len(lam) - 1
                for i in range(0, int((last + 1) / 2) + 1):
                    if not peq(padd(lam[i], lam[last - i]), 0):
                        message = "(" + ",".join(pstr(x) for x in lam)
                        message += (
                            ") is not Hermitian. \nIn type A every "
                            "coordinate must occur in a plus-minus pair, "
                            "such as (2,1,0,-1,-2).\nIf you enter all "
                            "non-negative coordinates it will be "
                            "symmetrized, \ne.g. (2,1,0) will be replaced "
                            "by  (2,1,0,-1,-2)\n\n")
                        return Result(type=type, lam=lam,
                                      lambdaWeight=lambdaWeight,
                                      lambdaWeightGap=lambdaWeightGap,
                                      result="stop", psd=0,
                                      message=message)
            rank = len(lam) - 1
        else:
            rank = len(lam)
            lam = dominant(type, lam)
        lambdaWeight = matrixMult([lam], transpose(simpleCoRoots(type, rank)))[0]
        lambdaWeightGap = matrixMult(
            [lam], transpose(simpleCoRoots(type, rank, "gap")))[0]

    lambdaStandardWeight = "(" + ",".join(pstr(x) for x in lam)
    lambdaStandardWeight += ")[" + ",".join(pstr(x) for x in lambdaWeight)
    lambdaStandardWeight += "]"
    if type == "D" and 2 * int(rank / 2) != rank and not peq(lam[-1], 0):
        message = "(" + ",".join(pstr(x) for x in lam)
        message += (") is not Hermitian.\nIn type D(odd) there must be a 0 "
                    "entry\n\n")
        return Result(type=type, lam=lam, lambdaWeight=lambdaWeight,
                      lambdaWeightGap=lambdaWeightGap, result="stop",
                      psd=0, message=message)

    integerClasses = sortIntegerClasses(lam)

    base = None
    if re.search("[CD]+", type):
        base = "0"
    elif re.search("B", type):
        base = "1/2"
    bkey = "" if base is None else base
    classicalPart = integerClasses.get(bkey) or []
    O_0, M_0, h_0, leftOver = getClassicalInfo(type, classicalPart)
    O_0, O_1, M_0, M_1, h_0, h_1, nu_0, nu_1 = moveToGL(O_0, M_0, h_0)
    integerClasses[bkey] = leftOver
    O_2, M_2, h_2, nu_2 = getGLInfo(type, integerClasses)
    O_GL = O_1 + O_2
    M_GL = M_1 + M_2
    h_GL = h_1 + h_2
    nu_GL = nu_1 + nu_2
    O = list(O_0) + flatten(O_GL)
    O = sort_desc(O)
    Z, nu = centralizer(type, O, O_0, O_GL, nu_GL)
    psd, result = test(type, Z, nu)

    if verbose > 1:
        p("\n\nG=")
        n = len(lam)
        if type == "A":
            p("GL(%d)" % n)
        elif type == "B":
            p("SO(%d)" % (2 * n + 1))
        elif type == "C":
            p("Sp(%d)" % (2 * n))
        elif type == "D":
            p("SO(%d)" % (2 * n))
        p(" (type %s%d)" % (type, n))
        p("\nlambda (std)[weight]:%s\n" % lambdaStandardWeight)
        p("\nOrbit O: (" + ",".join(pstr(x) for x in O) + ")\n")
        p("Centralizer Z= S[")
        keys = sorted(Z.keys(), key=lambda k: -pnum(k))
        p("x".join("%s(%s)" % (Z[j][1], Z[j][0]) for j in keys))
        p("]\n")
        output(type, O_0, M_0, h_0, O_GL, M_GL, h_GL, nu_GL, out)
        p("\nrow/mult/Z\tunitary\tnu\n")
        keys = list(psd.keys())
        for i in O_0:
            if not psd.get(_hkey(i)):
                keys.append(_hkey(i))
        keys = sorted(keys, key=lambda k: -pnum(k))
        for i in keys:
            if psd.get(i) is not None:
                mult = Z[i][0]
                factor = Z[i][1]
                p("/".join([i, pstr(mult), "%s(%s)" % (factor, mult)]))
                p("\t" + ("+" if psd[i] == 1 else "-"))
                p("\t(" + ",".join(pstr(x) for x in nu[i]) + ")\n")
            else:
                p("%s/1/O(1)\t+\t*\n" % i)
    if verbose > 0:
        p("%s: %s\n" % ("PASS" if result == 1 else "FAIL",
                        lambdaStandardWeight))
    if verbose > 1:
        p("------------------------------\n")

    return Result(type=type, lam=lam, lambdaWeight=lambdaWeight,
                  lambdaWeightGap=lambdaWeightGap, result=result, psd=psd,
                  message=message, Z=Z, O=O, O_0=O_0, M_0=M_0, h_0=h_0,
                  O_GL=O_GL, M_GL=M_GL, h_GL=h_GL, nu_GL=nu_GL, nu=nu)


def rhoGL(n):
    """rho for GL(n): ((n-1)/2, (n-3)/2, ..., -(n-1)/2)."""
    if n == 0:
        return [0]
    rv = []
    for j in range(n):
        entry = psub(frac(n - 1, 2), j)
        s = pstr(entry)
        if "/1" in s:
            entry = s.replace("/1", "", 1)
        rv.append(entry)
    return rv


def centralizer(type, O, O_0, O_GL, nu_GL):
    """Centralizer Z of the orbit O: row length -> [multiplicity, factor].

    Also returns tau: row length -> list of the nu's of the GL factors
    with that row length.
    """
    tau = {}
    for i in range(len(O_GL)):
        tau.setdefault(_hkey(O_GL[i][0]), []).append(nu_GL[i])
    Z = {}
    if not O:
        # Perl: undef..undef gives the single key "" (this happens in
        # type B when every coordinate is a half-integer that is not
        # absorbed by getClassicalInfo).
        Z[""] = [0, None]
    elif pnum(O[0]) <= pnum(O[-1]):
        # (Perl: foreach $i ($O->[0]..$O->[-1]) -- only non-empty when
        # all rows have the same length.)
        for i in range(int(pnum(O[0])), int(pnum(O[-1])) + 1):
            Z[_hkey(i)] = [0, None]
    for i in O:
        Z.setdefault(_hkey(i), [0, None])[0] += 1
    for i in Z:
        if type == "A":
            Z[i][1] = "GL"
        else:
            Z[i][1] = "O" if (int(pnum(i)) & 1) else "Sp"
    return Z, tau


def output(type, O_0, M_0, h_0, O_GL, M_GL, h_GL, nu_GL, out):
    """The M/O/h/nu table in the text (command line) output."""
    p = out.append
    indices = sorted(range(len(M_GL)),
                     key=_key(lambda a, b: pcmp(M_GL[b], M_GL[a])))
    p("\nM\t\tO\t\th\t\tnu\n")
    m0 = M_0[0] if M_0 else 0
    if type == "C":
        p("Sp(%s)\t\t" % pstr(2 * m0))
    elif type == "B":
        p("SO(%s)\t\t" % pstr(2 * m0 + 1))
    elif type == "D":
        p("SO(%s)\t\t" % pstr(2 * m0))
    p("(" + ",".join(pstr(x) for x in O_0) + ")\t\t")
    p("(" + ",".join(pstr(x) for x in flatten(h_0)) + ")\n")
    for i in indices:
        p("GL(%s)\t\t" % pstr(M_GL[i]))
        p("(" + ",".join(pstr(x) for x in O_GL[i]) + ")\t\t")
        p("(" + ",".join(pstr(x) for x in h_GL[i]) + ")\t\t")
        p(pstr(fix_str(nu_GL[i])) + "\n")


def flatten(arg):
    """[[2,3],[1,2,3],[1,2]] -> [2,3,1,2,3,1,2]."""
    rv = []
    for a in arg:
        rv.extend(a)
    return rv


def sortIntegerClasses(lam):
    """Group the coordinates of lambda by their class mod Z (up to sign).

    Returns a dict: remainder r (0 <= r <= 1/2, as a string key) ->
    list of the coordinates congruent to +-r mod Z.  Integer
    coordinates become plain integers, the others fractions.
    """
    h = {}
    for i in lam:
        s = pstr(i)
        if "/" in s:
            num, denom = (s.split("/") + [""])[:2]
            i = frac(num, denom)
        else:
            i = frac(i)
        j = i.plist(mixed=True)
        rem = frac(j[1], j[2])
        if peq(rem, 0):
            rem = 0
            i = j[0]
        rem = lessThanHalf(rem)
        if peq(rem, 0):
            rem = 0
        h.setdefault(_hkey(rem), []).append(i)
    return h


def getClassicalInfo(type, arg):
    """The distinguished (classical) part of the orbit.

    From the integral (C, D) or half-integral (B) coordinates build
    strings 0,1,..,k / 1,..,k / 1/2,..,k; each string is a row of the
    orbit O_0 and gives part of h.  Returns (O_0, [M_0], h_0, leftover).
    """
    if type == "A":
        return [], [], [], []
    classicalPart = sort_asc(list(arg))
    rvalues = {}
    for i in classicalPart:
        k = _hkey(i)
        rvalues[k] = rvalues.get(k, 0) + 1

    def rv(k):
        return rvalues.get(_hkey(k), 0)

    h = []
    O = []
    M = 0
    last = classicalPart[-1] if classicalPart else None
    if re.search("[CD]", type):
        for j in range(len(classicalPart)):
            parity = ((j + 1) & 1) if type == "D" else (j & 1)
            if rv(1 - parity) == 0:
                break
            string = []
            if j == 0 and not rv(1):
                continue
            for i in range(1 - parity, int(pnum(last)) + 1):
                if rv(i) == 0:
                    break
                string.append(i)
                rvalues[_hkey(i)] = rv(i) - 1
            string = sort_desc(string)
            h.append(string)
            O.append(2 * len(string) - 2 * parity + 1)
            M += len(string)
    else:
        for j in range(len(classicalPart)):
            if rv("1/2") == 0:
                break
            string = []
            largest = last
            for i in range(0, int(pnum(psub(largest, 0.5))) + 1):
                I = frac(i + 0.5)
                if rv(I) == 0:
                    break
                string.append(I)
                rvalues[_hkey(I)] = rv(I) - 1
            string = sort_desc(string)
            h.append(string)
            O.append(2 * len(string))
            M += len(string)

    leftOver = []
    top = int(pnum(last)) if classicalPart else -1
    for i in range(0, top + 1):
        for _ in range(int(pnum(rv(i)))):
            leftOver.append(i)

    Ms = [M]
    sizeO = sum(O)
    if type == "C" and not (sizeO & 1):
        O.append(1)
        sizeh = sum(len(x) for x in h)
        if sizeh < M:
            h.append([0])
    elif type == "D" and (sizeO & 1):
        O.append(1)
        h.append([0])
    return O, Ms, h, leftOver


def getGLInfo(type, hashes):
    """The GL factors of the Levi coming from the non-classical part.

    For each class of coordinates, repeatedly take the largest y and
    the longest string y, y-1, ..., (using |z| in types B, C, D); a
    string of length k gives GL(k) with h = rho(GL(k)) and
    nu = y - (k-1)/2.
    """
    O, M, h, nu = [], [], [], []
    for remainder in hashes:
        x = sort_desc(hashes[remainder])
        rvalues = {}
        for i in x:
            k = _hkey(i)
            rvalues[k] = rvalues.get(k, 0) + 1

        def rv(k):
            return rvalues.get(_hkey(k), 0)

        while True:
            y = None
            for z in x:
                if rv(z) > 0:
                    y = z
                    break
            if y is None or rv(y) == 0:
                break
            rvalues[_hkey(y)] = rv(y) - 1
            string = [y]
            top = int(pnum(padd(pmul(2, y), 1)))
            for i in range(1, top + 1):
                z = psub(y, i)
                if type == "A":
                    if rv(z):
                        string.append(z)
                        rvalues[_hkey(z)] = rv(z) - 1
                    else:
                        break
                else:
                    az = pabs(z)
                    if rv(az):
                        string.append(z)
                        rvalues[_hkey(az)] = rv(az) - 1
                    else:
                        break
            k = len(string)
            shift = psub(string[0], n_div(k - 1, 2))
            rho = []
            for i in range(k):
                a = n_div(k - 1, 2) - i
                if not (k & 1):
                    a = pstr(2 * a) + "/2"
                rho.append(a)
            if type == "A":
                O.append([k])
            else:
                O.append([k, k])
            M.append(k)
            nu.append(shift)
            h.append(rho)
    return O, M, h, nu


def dominant(type, arg):
    """Make lambda dominant: absolute values (not in type A), decreasing."""
    arg = list(arg)
    if type != "A":
        arg = [pabs(x) for x in arg]
    return sort_desc(arg)


def moveToGL(O_0, M_0, h_0):
    """Move pairs of equal rows of the classical orbit to GL factors."""
    O_0 = list(O_0)
    h_0 = list(h_0)
    M_0 = list(M_0)
    O_1, h_1, M_1, nu_1 = [], [], [], []
    while True:
        if not O_0:
            break
        found = False
        stop = False
        for i in range(len(O_0)):
            if i >= len(O_0) - 1:
                stop = True
                break
            if O_0[i] == O_0[i + 1]:
                rowLength = O_0[i + 1]
                del O_0[i:i + 2]
                del h_0[i:i + 2]
                M_0[0] = M_0[0] - rowLength
                O_1.append([rowLength, rowLength])
                h_1.append(rhoGL(rowLength))
                nu_1.append(0)
                M_1.append(rowLength)
                found = True
                break
        if stop:
            break
        if not found:
            break
    nu_0 = [[0]]
    if M_0 and M_0[0] == 0:
        h_0 = []
    return O_0, O_1, M_0, M_1, h_0, h_1, nu_0, nu_1


def test(type, Z, tau):
    """Test nu on each factor of Z; returns (psd, result).

    psd: row length -> 1 (unitary) or 0, for the rows carrying nu.

    Table implemented (row length, multiplicity m):
      G=A: Z-factor GL(m), dual type A(m-1)
      G=B: odd rows: C(m/2)->B;  even rows: m even D(m/2), m odd C((m-1)/2)
      G=C,D: even rows: B(m/2);  odd rows: m even D(m/2), m odd C((m-1)/2)
    """
    result = 1
    psd = {}
    for rowLength in Z:
        mult = Z[rowLength][0]
        rl = int(pnum(rowLength))
        if type == "A":
            rowType, rank = "A", mult - 1
        elif type == "B":
            if rl & 1:
                rowType, rank = "B", n_div(mult, 2)
            elif not (mult & 1):
                rowType, rank = "D", n_div(mult, 2)
            else:
                rowType, rank = "C", n_div(mult - 1, 2)
        else:
            if not (rl & 1):
                rowType, rank = "B", n_div(mult, 2)
            elif not (mult & 1):
                rowType, rank = "D", n_div(mult, 2)
            else:
                rowType, rank = "C", n_div(mult - 1, 2)
        if tau.get(rowLength):
            psd[rowLength] = testComplementarySeries(rowType, rank,
                                                     tau[rowLength])
        if not psd.get(rowLength) and tau.get(rowLength):
            result = 0
    return psd, result


def testComplementarySeries(rowType, rank, nu):
    """Is nu in the 0-complementary series of the given type and rank?"""
    nu = sort_desc(list(nu))
    if rowType == "D" and int(rank / 2) - n_div(rank, 2) != 0 \
            and not peq(nu[-1], 0):
        return 0
    if ple(nu[0], 0.5):
        return 1
    if pgt(nu[0], 1):
        return 0
    if re.search("A|B", rowType) and pgt(nu[0], 0.5):
        return 0
    for i in range(len(nu)):
        for j in range(i + 1, len(nu)):
            if peq(padd(nu[i], nu[j]), 1):
                return 0
    mu = sorted(nu, key=_key(lambda a, b: pcmp(lessThanHalf(a),
                                                lessThanHalf(b))))
    sums = []
    for i in range(len(mu)):
        s = 0
        if not pgt(mu[i], 0.5):
            continue
        for j in range(i + 1, len(mu)):
            if pgt(mu[j], 0.5):
                break
            s += 1
        sums.append(s)
    if sums and sums[-1] & 1:
        return 0
    for i in range(len(sums) - 1):
        if not (sums[i] & 1):
            return 0
    return 1


def lessThanHalf(a):
    """Replace 1/2 < x < 1 by 1 - x."""
    return a if ple(frac(a), 0.5) else psub(1, a)


def partitions(n):
    """All partitions of n (lists), in the order of the Perl program."""
    if n - 1 == 0:
        return [[1]]
    ps = []
    for x in partitions(n - 1):
        ps.append(x + [1])
        if len(x) == 1 or x[-2] > x[-1]:
            y = list(x)
            y[-1] += 1
            ps.append(y)
    return ps


HELP_TEXT = """
Unitarity tester for Classical Groups. Given a type A, B, C or D and a
list lambda of real numbers (standard coordinates), decide whether the
spherical representation X(lambda) of the corresponding split real or
p-adic group is unitary, and give some other information about it.
This implements an algorithm due to Dan Barbasch.
"""