#!/usr/bin/python
# -*- coding: utf-8 -*-

r"""
Python script to create the Markov trace related source for the
``database_cubic_hecke`` package in ``py`` format in the given path.
This utility should be used to switch to a new version of the data
files if the corresponding files have run out of date.

.. NOTE::

    This function needs Sage together with the code of Trac ticket
    :trac:`29717`.

INPUT:

- ``path`` -- name of path where to store the generated  files
  (if not provided, the the subdirectory ``database_cubic_hecke``)
  under the current working directory is used)

EXAMPLES::

    sage create_markov_trace_data.py
    sage create_markov_trace_data.py ~

Be aware that these invocations will run for a couple of hours!
"""

##############################################################################
#       Copyright (C) 2022 Sebastian Oehms <seb.oehms@gmail.com>
#
# This program is free software: you can redistribute it and/or modify
# it under the terms of the GNU General Public License as published by
# the Free Software Foundation, either version 2 of the License, or
# (at your option) any later version.
#                  http://www.gnu.org/licenses/
##############################################################################

import sys
from os import environ
from os.path import join
from logging import info, basicConfig
from database_cubic_hecke import FileNames
from database_cubic_hecke import __version__ as this_version
from sage.all_cmdline import *



basicConfig(format='%(asctime)s %(message)s', level=20)

cmdline_args = sys.argv[1:]

path = None

if len(cmdline_args) > 0:
    path = cmdline_args[0]
    if not exists(path):
        raise ValueError("the path '%s' does not exist" % path)

if not path:
    pwd = environ['PWD']
    path = join(pwd, 'database_cubic_hecke')


# -----------------------------------------------------------------------------
# Templates
# -----------------------------------------------------------------------------
py_header_text = r"""# ------------------------------------------------------------
# This file was generated using create_markov_trace_data.py
#
# on    : Version %s
# under : Sage %s
# using : Python %s
# ------------------------------------------------------------

"""
function_name = 'read_%s'
template=r"""# generated data function for cubic Hecke algebra
def %s(bas_ele, variables, num_strands=4):
    %s
    %s = variables
    data = {}
    data[2] = %s
    data[3] = %s
    data[4] = %s
    return data[num_strands][bas_ele]

"""

doc=r"""{}
    Return precomputed Markov trace coefficients.

    This code has been generated by ``create_markov_trace_data.py`` (from
    the `database_cubic_hecke repository  <https://github.com/soehms/database_cubic_hecke>`__),
    please don't edit.

    INPUT:

    - ``bas_ele`` -- instance of enum :class:`MarkovTraceModuleBasis`
    - ``variables`` -- tuple consisting of the variables used in
      the coefficients
    - ``num_strands`` -- integer indicating the number of strands (default 4)


    OUTPUT:

    A list of the coefficients. The i'th member corresponds to the i'th%s
    {}
""".format('r"""', '"""')

doc_reg=r"""
    basis element.

    EXAMPLES::

        >>> from database_cubic_hecke.markov_trace_coeffs import read_markov
        >>> from sympy import var
        >>> u, v, w, s = var('u, v, w, s')
        >>> variables = (u, v, w, s)
        >>> read_markov('U2', variables, num_strands=3) # doctest: +NORMALIZE_WHITESPACE
        [0, s, 1/s, s, 1/s, 0, 0, 0, 0, -s*v, s, s, -s*u/w, -v/s, 1/s,
         0, 0, 0, 0, 1/s, -u/(s*w), -v/s, 0, 0]
"""


doc_irr=r"""
    irreducible character

    EXAMPLES::

        >>> from database_cubic_hecke.markov_trace_coeffs_irr import read_markov_irr
        >>> from sympy import var, I, sqrt
        >>> zeta3 = (sqrt(3)*I-1)/2
        >>> a, b, c, s = var('a, b, c, s')
        >>> variables = (a, b, c, s, zeta3)
        >>> read_markov_irr('U1', variables, num_strands=2) # doctest: +NORMALIZE_WHITESPACE
        [(a*b*c + a*s**2)/(a**2*s - a*b*s - a*c*s + b*c*s),
         (-a*b*c - b*s**2)/(a*b*s - a*c*s - b**2*s + b*c*s),
         (a*b*c + c*s**2)/(a*b*s - a*c*s - b*c*s + c**2*s)]
"""


# -----------------------------------------------------------------------------
# Auxillary function to generate the function code
# -----------------------------------------------------------------------------
def _write_data_function(filename, data):
    r"""
    Generate code for functions to access the data.
    """
    from textwrap import fill
    def fl(s):
        return fill(str(s), subsequent_indent='        ')
    func_name = function_name %filename.cache()

    if filename == FileNames.markov_coeffs_irr:
        var = 'a, b, c, s, zeta3'
        varexp = '3, 5, 7, 11, 13'
        docstr  = doc %doc_irr
    else:
        var = 'u, v, w, s'
        varexp = '3, 5, 7, 13'
        docstr  = doc %doc_reg

    fldata = [fl(str(d).replace('^', '**')) for d in data]
    d2, d3, d4 = fldata
    func  = template %(func_name, docstr, var, d2, d3, d4)

    header = py_header_text %(this_version.value, version(), sys.version.split(' ')[0])

    import_filen = join(path, filename.py())
    with open(import_filen, 'w') as f:
        f.write(header)
        f.write(func)


# --------------------------------------------------------------------------
# calc_markov_trace_coeffs
# --------------------------------------------------------------------------
def calc_markov_trace_coeffs(CHA, bas_ele=None, mcfs_irr=None):
    r"""
    Return a list of formal Markov traces of the basis elements of ``CHA``.

    INPUT:

    - ``CHA`` -- instance of :class:`CubicHeckeAlgebra`
    - ``bas_ele`` -- basis element of the Markov trace module of ``CHA``
      (see :meth:`_markov_trace_module`). If given then the list of
      coefficients for the basis elements of ``CHA`` with respect to
      ``bas_ele`` is returned. Else the entries of the list are elements
      of the Markov trace module.
    - ``mcfs_irr`` -- (optional) Markov trace coefficients of the irreducible
      characters of ``CHA`` as returned by :func:`calc_markov_trace_irr_coeffs`


    OUTPUT:

    A list of elements of the Markov trace module (over the Markov trace
    version of the generic base ring). Each entry of the list corresponds
    to the according basis element of ``CHA``. If ``bas_ele`` is given
    then instead of the module element its coefficients with respect to
    the given ``bas_ele`` is returned.

    EXAMPLES::

        sage: CHA2 = algebras.CubicHecke(2)
        sage: calc_markov_trace_coeffs(CHA2)
        [B[U2], s*B[U1], 1/s*B[U1]]
        sage: M = _[0].parent(); M
        Free module generated by {U1, U2}
          over Multivariate Polynomial Ring in u, v, w, s
          over Integer Ring localized at (s, w, v, u)
        sage: U1, U2 = M.basis().keys()
        sage: calc_markov_trace_coeffs(CHA2, bas_ele=U1)
        [0, s, 1/s]
        sage: calc_markov_trace_coeffs(CHA2, bas_ele=U2)
        [1, 0, 0]
    """
    n = CHA.strands()

    # ----------------------------------------------------------------------
    # Markov trace module
    # ----------------------------------------------------------------------
    MTM = CHA._markov_trace_module()
    MTMbas = MTM.basis().keys()
    L = MTM.base_ring()

    # ----------------------------------------------------------------------
    # Setting up filecache handling
    # ----------------------------------------------------------------------
    last_bas_ele = MTMbas[len(MTMbas)-1]

    # ----------------------------------------------------------------------
    # If no bas_ele is given proceed by recursion
    # ----------------------------------------------------------------------
    if not bas_ele:
        return {bas_ele: calc_markov_trace_coeffs(CHA, bas_ele=bas_ele, mcfs_irr=mcfs_irr) for bas_ele in MTMbas}

    O = CHA.get_order()
    l = len(O)

    info('Preparing calculation of Markov trace coefficients for %s on %s strands' %(bas_ele.name, n))

    u, v, w, s = L.gens()

    # ----------------------------------------------------------------------
    # If precalculated coeffocients of the irreducible modules exist in 
    # filcache take it from there. Otherwise calculate it in 
    # CHA._markov_trace_irr_coeffs
    # ----------------------------------------------------------------------
    from database_cubic_hecke.markov_trace_coeffs_irr import read_markov_irr

    F = CHA._markov_trace_module(extended=True, field_embedding=True).base_ring()
    a, b, c, si, = F.gens()
    zeta3 = F.base_ring().gen()
    variables = (a, b, c, si, zeta3)
    if mcfs_irr:
        irr_coeff = mcfs_irr[bas_ele]
    else:
        irr_coeff = read_markov_irr(bas_ele.name, variables, num_strands=n)
    cf_vect = vector(irr_coeff)

    # ----------------------------------------------------------------------
    # Setup conversion of the coefficients to L
    # ----------------------------------------------------------------------
    P = L.extension_ring()
    Ps = P.as_splitting_algebra()
    EZ = ZZ[P.variable_names()]
    embEZ = EZ.hom(P.gens())
    chars = CHA.characters()
    num_char = len(chars)
    def mtr(ele):
        ch = vector([chars[j](ele) for j in range(num_char)])
        return cf_vect * ch

    subs_dict = L.gens_dict_recursive()
    from sage.misc.sage_eval import sage_eval

    def convert_coeff(g):
        cf = mtr(CHA(g))
        num = cf.numerator()
        den = cf.denominator()
        num_Ps = Ps(embEZ(EZ(num.dict())))
        den_Ps = Ps(embEZ(EZ(den.dict())))
        num_L = sage_eval(str(num_Ps), locals=subs_dict)
        den_L = sage_eval(str(den_Ps), locals=subs_dict)
        res = L(num_L/den_L)
        return res

    info('Start calculating %s (%s strands)' %(bas_ele.name, n))

    # ----------------------------------------------------------------------
    # Perform the calculation of the coefficients 
    # ----------------------------------------------------------------------
    mtcf = []
    for i in range(l):
        info('    Calculate %sth coefficient (of %s)  for %s on %s strands)' %(i, O[i], bas_ele.name, n))
        cf = convert_coeff(O[i])
        mtcf.append(cf)

    info('Calculation of Markov trace coefficients for %s finished (%s strands)' %(bas_ele.name, n))
    return mtcf


def calc_markov_trace_irr_coeffs(CHA, mcfs_sub=None):
    r"""
    Return a list of formal Markov traces as linear forms on the irreducible
    characters of ``CHA``.

    INPUT:

    - ``CHA`` -- instance of :class:`CubicHeckeAlgebra`
    - ``mcfs_sub`` -- (optional) Markov trace coefficients of the cubic Hecke subalgebra
      as returned by :func:`calc_markov_trace_coeffs`

    OUTPUT:

    A dictionary from the irreducible representations of ``CHA`` to elements
    of the Markov trace module (over the field embedding of the Markov trace
    version of the generic extension ring).

    EXAMPLES::

        sage: CHA2 = algebras.CubicHecke(2)
        sage: calc_markov_trace_irr_coeffs(CHA2)
        {U1: ((a*b*c + a*s^2)/(a^2*s - a*b*s - a*c*s + b*c*s),
              (-a*b*c - b*s^2)/(a*b*s - b^2*s - a*c*s + b*c*s),
              (a*b*c + c*s^2)/(a*b*s - a*c*s - b*c*s + c^2*s)),
         U2: ((-a*b - a*c)/(a^2 - a*b - a*c + b*c),
              (a*b + b*c)/(a*b - b^2 - a*c + b*c),
              (-a*c - b*c)/(a*b - a*c - b*c + c^2))}
    """

    # --------------------------------------------------------
    # Preparations
    # --------------------------------------------------------
    n = CHA.strands()
    chars = CHA.characters()
    num_char = len(chars)
    info('Preparing calculation of irreducible Markov trace coefficients for %s strands' %n)

    # --------------------------------------------------------
    # Define modules and rings
    # --------------------------------------------------------
    MTM    = CHA._markov_trace_module(extended=True, field_embedding=True)
    MTMbas = MTM.basis().keys()
    F = MTM.base_ring()
    a, b, c, s = F.gens()

    from sage.matrix.constructor import matrix

    if CHA.strands() == 2:
        U1, U2 = MTMbas
        eq = matrix(F, 3, lambda i, j: chars[j](CHA(CHA.get_order()[i])))
        eqi = ~eq
        v1 = vector(F, [0, s, ~s])
        v2 = vector(F, [1, 0, 0])
        return {U1: eqi*v1, U2: eqi*v2}

    MTMbas_new = [b for b in MTMbas if b.strands() == n]

    # --------------------------------------------------------
    # incrediences from the sub algebra on one strand less
    # --------------------------------------------------------
    sub = CHA.cubic_hecke_subalgebra()
    sub_basis = [sub(g) for g in sub.get_order()]
    sub_dim = sub.dimension()
    sub_MTM    = sub._markov_trace_module(extended=True, field_embedding=True)
    sub_MTMbas = sub_MTM.basis().keys()
    if mcfs_sub:
        # Intermediate ring Q is needed (see :meth:`formal_markov_trace`
        # of :class:`CubicHeckeElement`)
        Q = sub._markov_trace_module(extended=True, field_embedding=False).base_ring()
        def to_mtr(i):
            return sum(F(Q(mcfs_sub[k][i]))*sub_MTM(k) for k in sub_MTMbas)
        sub_mtr = [to_mtr(i) for i in range(sub_dim)]
    else:
        sub_mtr = [b.formal_markov_trace(extended=True, field_embedding=True) for b in sub_basis]
    new_basis_ele = [CHA(b.braid_tietze()) for b in MTMbas_new]

    # --------------------------------------------------------
    # setting up equation
    # --------------------------------------------------------
    info('Setting up equation (%s strands)' %n)

    g = CHA.gen(CHA.ngens()-1)
    eq_p = matrix(F, sub_dim, num_char, lambda i,j: chars[j](CHA(sub_basis[i])*g))
    eq_n = matrix(F, sub_dim, num_char, lambda i,j: chars[j](CHA(sub_basis[i])*~g))
    eq_b = eq_p.stack(eq_n)

    # --------------------------------------------------------
    # solve the equation
    # --------------------------------------------------------
    info('Calculate kernel(%s strands)' %n)
    ker = eq_b.right_kernel()

    dk = ker.dimension()
    kerm = ker.basis_matrix()

    # --------------------------------------------------------
    # adjusting kernel to new variables
    # --------------------------------------------------------
    info('Calculate adjusting equation(%s strands)' %n)
    def vectchars(ele):
        return vector(F, [chars[j](ele) for j in range(num_char)])

    eq_adjust = matrix(F, dk, dk, [list(kerm*vectchars(ele)) for ele in new_basis_ele])

    def get_irr_coeff(bas_ele):
        pre_image = vector(F, [0]*num_char)
        image_new_basis = vector(F, [0]*dk)
        claim_new_basis = vector(F, [0]*dk)
        if bas_ele in sub_MTMbas:
            sub_mtr_b = vector(F, [s*mtr.coefficient(bas_ele) for mtr in sub_mtr] + [~s*mtr.coefficient(bas_ele) for mtr in sub_mtr] )
            info('   Solving preimage for adjusting equation for %s' %bas_ele.name)
            pre_image = eq_b.solve_right(sub_mtr_b)
            image_new_basis = vector(F, [pre_image*vectchars(ele) for ele in new_basis_ele])
        else:
            claim_new_basis = vector(F, [MTMbas_new.index(bas_ele) == i for i in range(dk)])
        info('   Solving adjusting equation for %s' %bas_ele.name)
        cfs_adjust_kernel_basis = eq_adjust.solve_right(claim_new_basis - image_new_basis)
        return pre_image + cfs_adjust_kernel_basis*kerm

    info('Solve adjusting equations (%s strands)' %n)

    irr_coeff = {bas_ele: get_irr_coeff(bas_ele) for bas_ele in MTMbas}

    # --------------------------------------------------------
    # simplification
    # --------------------------------------------------------
    from sage.misc.misc_c import prod

    def red_coeff(cf):
        cfn = cf.numerator().factor()
        cfd = cf.denominator().factor()
        un = cfn.unit()
        ud = cfd.unit()
        prn = prod(f**e for (f,e) in cfn)
        prd = prod(f**e for (f,e) in cfd)
        u = un/ud
        return u*prn/prd

    info('Reduce adjusted solution(%s strands)' %n)

    irr_coeff_red = {b: vector([red_coeff(cf) for cf in irr_coeff[b]]) for b in irr_coeff.keys()}

    irr_coeff = irr_coeff_red
    info('Calculation of irreducible Markov trace coefficients finished (%s strands)' %n)
    return irr_coeff

def create_markov_trace_coeffs():
    r"""
    Generate Python code to obtain formal Markov traces coefficients on the irreducible
    characters of the cubic Hecke algebras.
    """
    CHA4 = algebras.CubicHecke(4)
    CHA3 = CHA4.cubic_hecke_subalgebra()
    CHA2 = CHA3.cubic_hecke_subalgebra()
    def prepare_data(data):
        return {k.name:[cf for cf in v] for k, v in data.items()}

    irr_coeffs2 = calc_markov_trace_irr_coeffs(CHA2)
    coeffs2     = calc_markov_trace_coeffs(CHA2, mcfs_irr=irr_coeffs2)

    irr_coeffs3 = calc_markov_trace_irr_coeffs(CHA3, mcfs_sub=coeffs2)
    coeffs3     = calc_markov_trace_coeffs(CHA3, mcfs_irr=irr_coeffs3)

    irr_coeffs4 = calc_markov_trace_irr_coeffs(CHA4, mcfs_sub=coeffs3)
    coeffs4     = calc_markov_trace_coeffs(CHA4, mcfs_irr=irr_coeffs4)

    fname = FileNames.markov_coeffs_irr
    data = (prepare_data(cf) for cf in [irr_coeffs2, irr_coeffs3, irr_coeffs4])
    _write_data_function(fname, data)

    fname = FileNames.markov_coeffs
    data = (prepare_data(cf) for cf in [coeffs2, coeffs3, coeffs4])
    _write_data_function(fname, data)


create_markov_trace_coeffs()
