{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "KL->Standard completed in 7.152557373046875e-07 seconds.\n",
      "Success! Rewrite took 0.615220308303833 seconds.\n",
      "1*N[1, 2, 0, 2, 1, 2, 1, 0] + v^2*N[1, 2, 1, 0]\n"
     ]
    }
   ],
   "source": [
    "## Companion code for the article, \"Kazhdan-Lusztig polynomials for \\tilde{B}_2\" by Karina Batistelli, Aram Bingham, and David Plaza\n",
    "\n",
    "import time\n",
    "\n",
    "# First we devine the polynomial ring where KL polynomials live, \\Z[v,v^{-1}]\n",
    "R.<v> = LaurentPolynomialRing(ZZ)\n",
    "# The IwahoriHeckeAlgebra function will provide us the generators and relations for affine type B2\n",
    "H = IwahoriHeckeAlgebra(['B',2,1], -v,v^(-1))  # the 1 after 'B', 2 indicates affine type\n",
    "# The -v, v^{-1} specifies the quadratic relation our realization obeys.\n",
    "# For T_s in the standard basis, T_s^2=(v^{-1}-v)T_s+1, as per \"Intro to Soergel bimodules\" (Elias et. al) (3.1)\n",
    "# We rename the basis elements for convenience. \n",
    "# T: standard basis, C: canonical basis, Cp: the \"other\" self dual basis from [KL79] which we will not use\n",
    "# C(x) corresponds to the Kazhdan-Lusztig basis as presented in Elias et. al\n",
    "T=H.T(); C=H.C(); Cp=H.Cp()\n",
    "# Affine Weyl group, the prefix means the default generators will be called s_0, s_1, s_2\n",
    "W = WeylGroup(['B', 2, 1],prefix=\"s\") \n",
    "[s0,s1,s2]= W.simple_reflections() # Let's shorten that to not have to write underscores\n",
    "\n",
    "# Now we define the elements N_x from the article\n",
    "def N(x):\n",
    "    return sum([v**(x.length()-y.length())*T[y] for y in W.bruhat_interval(1,x)])\n",
    "\n",
    "### now make the rewrite function\n",
    "# We will use the articles notation where H_x is the standard basis element that Sage denotes as something like T[x]\n",
    "def rwNH(h):    #rewrite the Hecke alg. elt. h in terms of N's and H's\n",
    "    start=time.time()  # This is a clumsy way to keep track of time, but it works.\n",
    "    hi=T(h)   #hi is for \"h in\" -- first convert h to standard basis, and we'll act on it peeling off N's\n",
    "    step1=time.time()\n",
    "    print('KL->Standard completed in '+str(step1-start)+' seconds.')\n",
    "    expansion=[]  # where we will collect terms of expansion  \n",
    "    while hi!=0: \n",
    "        terms=[str(w) for w in hi.terms()] # terms in sage's default order which is sort of by length but with some irregularities\n",
    "        sup=[]   #support, will be list of elements of the form T[1,2,0], e.g in decreasing length order\n",
    "        for x in terms:\n",
    "            init=\"T\"    \n",
    "            sup.append(x[x.find(init):])  # This is a goofy way to strip the coefficients: \n",
    "            # We find where T appears in the stringified term and take everything including and after that\n",
    "            # \n",
    "        sup.sort(key=len, reverse=True) # arranges T basis elts in decreasing length order (string length order and group length are compatible)\n",
    "        #print(terms) \n",
    "        #print(sup) # for debugging\n",
    "        \n",
    "        \n",
    "        index = [idx for idx, s in enumerate(terms) if sup[0] in s][0] # finds the list index of the longest elt in sages ordering\n",
    "        ltb= sage_eval(sup[0], locals={'T':T}) # converts string back to algebra element, the longest in sup (lt is for \"leading term\")\n",
    "        ltc=sage_eval(terms[index], locals={'T':T, 'v':v}).coefficients()[0] # now find coefficient of leading term\n",
    "        #print(ltb)\n",
    "        #print(hi.terms())\n",
    "        #print(ltc)\n",
    "        \n",
    "        rwtime0=time.time()\n",
    "        ltbrw = ltb.support()[0].reduced_word() # need lt underlying group elt as seq of coxeter gens)\n",
    "        ltw=W.from_reduced_word(ltbrw) # this is the actual elt in the gp, not the expression\n",
    "        rwtime1=time.time()\n",
    "        #print('Reduced word extraction took '+str(rwtime1-rwtime0)+' seconds.')\n",
    "        #ltc = lt.coefficients()[0]   #c is for 'coeff' \n",
    "        #print(ltc)\n",
    "        \n",
    "        ntime0=time.time()\n",
    "        Nlt=N(ltw)  # take the N for the leading term\n",
    "        ntime1=time.time()\n",
    "        #print('N construction took '+str(ntime1-ntime0)+' seconds.') # in practice this took very little time, just wanted to see what were the costly steps\n",
    "        #print(Nlt)\n",
    "        \n",
    "        stime0=time.time()\n",
    "        if \"-\" not in str(hi-ltc*Nlt): #as long as we produce no negative symbols....\n",
    "            hi -= ltc*Nlt   # subtract the N term (with coefficient) corresponding to leading term elt from the input\n",
    "            expansion.append(str(ltc)+'*N'+str(ltbrw)) # keep track of this expansion\n",
    "            #print(hi,\"Subtracted N\")\n",
    "        else:\n",
    "            hi -= ltc*ltb  # else we just subtract an H(x) with its coefficient and move on\n",
    "            expansion.append(str(ltc)+'*T'+str(ltbrw))\n",
    "            #print(hi,\"Subtracted T\")\n",
    "        stime1=time.time()\n",
    "       # print('Subtraction took '+str(stime1-stime0)+' seconds.')\n",
    "        \n",
    "    if hi==0:\n",
    "        end=time.time()\n",
    "        print('Success! Rewrite took '+ str(end-step1)+' seconds.' )\n",
    "    else:\n",
    "        print('Algo anda mal :(')\n",
    "    print(*expansion, sep=' + ')\n",
    "    \n",
    "rwNH(N(s1*s2*s0*s2*s1*s2*s1)*C[0])\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "KL->Standard completed in 2.1457672119140625e-06 seconds.\n",
      "Success! Rewrite took 3.8850955963134766 seconds.\n",
      "1*N[1, 2, 0, 2, 1, 2, 1, 0, 2, 1, 0] + v*N[1, 2, 0, 2, 1, 2, 1, 0]\n"
     ]
    }
   ],
   "source": [
    "rwNH(N(s1*s2*s0*s2*s1*s2*s1*s0*s2*s0)*C[1])\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "KL->Standard completed in 4.76837158203125e-07 seconds.\n",
      "Success! Rewrite took 0.11400485038757324 seconds.\n",
      "1*N[0, 2, 1, 2, 1]\n",
      "KL->Standard completed in 1.1920928955078125e-06 seconds.\n",
      "Success! Rewrite took 0.9381084442138672 seconds.\n",
      "1*N[0, 2, 1, 2, 1, 0, 2, 0] + v*N[1, 2, 0, 2, 0]\n",
      "KL->Standard completed in 1.6689300537109375e-06 seconds.\n",
      "Success! Rewrite took 1.6780059337615967 seconds.\n",
      "1*N[0, 2, 1, 2, 1, 0, 2, 1, 0, 2, 1] + v*N[1, 2, 0, 2, 1, 0, 2, 1] + v^2*N[0, 2, 1, 2, 1]\n",
      "KL->Standard completed in 1.430511474609375e-06 seconds.\n",
      "Success! Rewrite took 3.894035577774048 seconds.\n",
      "1*N[0, 2, 1, 2, 1, 0, 2, 1, 0, 2, 1, 0, 2, 0] + v*N[1, 2, 0, 2, 1, 0, 2, 1, 0, 2, 0] + v^2*N[0, 2, 1, 2, 1, 0, 2, 0] + v^3*N[1, 2, 0, 2, 0]\n",
      "KL->Standard completed in 2.86102294921875e-06 seconds.\n",
      "Success! Rewrite took 8.48439908027649 seconds.\n",
      "1*N[0, 2, 1, 2, 1, 0, 2, 1, 0, 2, 1, 0, 2, 1, 0, 2, 1] + v*N[1, 2, 0, 2, 1, 0, 2, 1, 0, 2, 1, 0, 2, 1] + v^2*N[0, 2, 1, 2, 1, 0, 2, 1, 0, 2, 1] + v^3*N[1, 2, 0, 2, 1, 0, 2, 1] + v^4*N[0, 2, 1, 2, 1]\n",
      "KL->Standard completed in 1.6689300537109375e-06 seconds.\n",
      "Success! Rewrite took 15.33332085609436 seconds.\n",
      "1*N[0, 2, 1, 2, 1, 0, 2, 1, 0, 2, 1, 0, 2, 1, 0, 2, 1, 0, 2, 0] + v*N[1, 2, 0, 2, 1, 0, 2, 1, 0, 2, 1, 0, 2, 1, 0, 2, 0] + v^2*N[0, 2, 1, 2, 1, 0, 2, 1, 0, 2, 1, 0, 2, 0] + v^3*N[1, 2, 0, 2, 1, 0, 2, 1, 0, 2, 0] + v^4*N[0, 2, 1, 2, 1, 0, 2, 0] + v^5*N[1, 2, 0, 2, 0]\n"
     ]
    }
   ],
   "source": [
    "# If one wants to rewrite Hecke alg. elements only in terms of N's, the code below will work\n",
    "# Note! These expansions are not unique! The algorithm given is greedy, so the output will depend on the order in which \n",
    "# the terms of standard basis expansion is fed to it.\n",
    "\n",
    "import time\n",
    "\n",
    "\n",
    "R.<v> = LaurentPolynomialRing(ZZ)\n",
    "H = IwahoriHeckeAlgebra(['B',2,1], -v,v^(-1))\n",
    "T=H.T(); Cp=H.Cp(); C=H.C()\n",
    "W = WeylGroup(['B', 2, 1],prefix=\"s\")\n",
    "[s0,s1,s2]= W.simple_reflections()\n",
    "\n",
    "# Same starting setup\n",
    "def N(x):\n",
    "    return sum([v**(x.length()-y.length())*T[y] for y in W.bruhat_interval(1,x)])\n",
    "\n",
    "def rwN(h):    #rewrite in terms of N's\n",
    "    start=time.time()\n",
    "    hi=T(h)   #hi is for \"h in\" -- we'll act on it peeling off N's, but first we convert to standard basis\n",
    "    step1=time.time()\n",
    "    print('KL->Standard completed in '+str(step1-start)+' seconds.')\n",
    "    expansion=[]  # where we will collect terms of expansion  \n",
    "    while hi!=0:\n",
    "        terms=[str(w) for w in hi.terms()] # terms in sage's default order\n",
    "        sup=[]   # \"support,\" will be list of elements of the form T[1,2,0], e.g in increasing length order\n",
    "        for x in terms:\n",
    "            start=\"T\"\n",
    "            sup.append(x[x.find(start):])\n",
    "        sup.sort(key=len, reverse=True) # arranges T basis elts in increasing length order\n",
    "        #print(terms)\n",
    "        #print(sup)\n",
    "        #print(sup[0])\n",
    "        \n",
    "        index = [idx for idx, s in enumerate(terms) if sup[0] in s][0] # finds the index of the longest elt in sages ordering\n",
    "        ltb= sage_eval(sup[0], locals={'T':T}) # converts string back to algebra element, s[0] is the longest in sup\n",
    "        ltc=sage_eval(terms[index], locals={'T':T, 'v':v}).coefficients()[0] # now find coefficient of leading term\n",
    "        #print(ltb)\n",
    "        #print(hi.terms())\n",
    "        #print(ltc)\n",
    "        \n",
    "        rwtime0=time.time()\n",
    "        ltbrw = ltb.support()[0].reduced_word() # need lt underlying group elt as seq of coxeter gens)\n",
    "        ltw=W.from_reduced_word(ltbrw) # this is the actual elt in the gp, not the expression\n",
    "        rwtime1=time.time()\n",
    "        #print('Reduced word extraction took '+str(rwtime1-rwtime0)+' seconds.')\n",
    "        #ltc = lt.coefficients()[0]   #c is for 'coeff' (lists are decreasing in length/bruhat it seems)\n",
    "        #print(ltc)\n",
    "        \n",
    "        ntime0=time.time()\n",
    "        Nlt=N(ltw)  # take the N for the leading term\n",
    "        ntime1=time.time()\n",
    "        #print('N construction took '+str(ntime1-ntime0)+' seconds.')\n",
    "        #print(Nlt)\n",
    "        \n",
    "        stime0=time.time()\n",
    "        hi -= ltc*Nlt   # and subtract it from the input\n",
    "        expansion.append(str(ltc)+'*N'+str(ltbrw))\n",
    "        stime1=time.time()\n",
    "        #print('Subtraction took '+str(stime1-stime0)+' seconds.')\n",
    "        \n",
    "    end=time.time()\n",
    "    print('Success! Rewrite took '+ str(end-step1)+' seconds.' )\n",
    "    print(*expansion, sep=' + ')    \n",
    "\n",
    "# We also define the elements D_x^y as defined in Section 6 of the article\n",
    "def D(x,y):\n",
    "    return sum([v**(x.length()-z.length())*T[z] for z in W.bruhat_interval(1,x) if not z.bruhat_le(y)]) \n",
    "\n",
    "# This function defines the \"prime\" involution that swaps s1 and s0\n",
    "def fi(x):\n",
    "    s=list(str(x)) # takes an element of the form e.g. s2*s1*s0*s2*s1*s2*s1*s0*s1\n",
    "    t=['0' if y=='1' else '1' if y=='0' else y for y in s] # once symbols are in a list, we swap 1's and 0's\n",
    "    t=sage_eval(''.join(t), locals={'s1':s1, 's2':s2, 's0':s0, 'C':C, 'T':T}) #recover the associated element\n",
    "    return t \n",
    "\n",
    "def Dtil(x): # \\tilde{D}(x) is D_x^{x'}, which is used in the definition of U_x\n",
    "    return D(x,fi(x))\n",
    "\n",
    "# U_x from Section 6\n",
    "def U(x):\n",
    "    return N(x)+D(fi(x),x)\n",
    "\n",
    "# For x of length one, the associated canonical basis element has this form; see Elias et al, ch. 3\n",
    "def Hb(x):\n",
    "    return T[x]+v \n",
    "\n",
    "# Elements x_n of the thick region as defined in the paper\n",
    "def x(n):\n",
    "    frag=[s1,s2,s1,s0,s2,s0] #repeating subsequence\n",
    "    reps=ceil(n/6)\n",
    "    enough=frag*reps #extend the repeating segment far enough\n",
    "    return prod(i for i in enough[:n]) # then take the first n symbols and multiply\n",
    "\n",
    "# x'_n, could have been defined using fi(x) but instead we define analagous to x(n)\n",
    "def xp(n):\n",
    "    frag=[s0,s2,s0,s1,s2,s1]\n",
    "    reps=ceil(n/6)\n",
    "    enough=frag*reps\n",
    "    return prod(i for i in enough[:n])\n",
    "\n",
    "# Elements e_n and e'_n of the other branches of the thick region\n",
    "def e(n):\n",
    "    return s1*xp(n)\n",
    "\n",
    "def ep(n):\n",
    "    return s0*x(n)\n",
    "\n",
    "# Elements d_n of the thin region as defined in the paper\n",
    "def d(n):\n",
    "    frag=[s2,s1,s2,s0]\n",
    "    reps=ceil(n/4)\n",
    "    enough=frag*reps\n",
    "    return prod(i for i in enough[:n])\n",
    "\n",
    "# d'_n\n",
    "def dp(n):     \n",
    "    frag=[s2,s0,s2,s1]\n",
    "    reps=ceil(n/4)\n",
    "    enough=frag*reps\n",
    "    return prod(i for i in enough[:n])\n",
    "\n",
    "# theta's of the big region as defined in the paper\n",
    "def theta(m,n):\n",
    "    if m%2==0:\n",
    "        return x(3*(m+1))*dp(4*n+1)\n",
    "    else:\n",
    "        return x(3*(m+1))*d(4*n+1)\n",
    "\n",
    "# \\theta'(m,n) (apply the prime involution)\n",
    "def thetap(m,n):\n",
    "    if m%2==0:\n",
    "        return xp(3*(m+1))*d(4*n+1)\n",
    "    else:\n",
    "        return xp(3*(m+1))*dp(4*n+1)    \n",
    "\n",
    "# According to results of the paper, we can define the canconical basis elements associated to theta's and \"friends\"\n",
    "# directly using the N elements. Hb is for \"H bar\", the notation used in the paper.\n",
    "def HbTheta(m): # This is for \\theta(m,0)\n",
    "    return sum([v^(2*i)*N(theta(m-2*i,0)) for i in range(floor(m/2)+1)])\n",
    "\n",
    "# H bar for theta prime\n",
    "def HbThetap(m): #\\theta'(m,0)\n",
    "    return sum([v^(2*i)*N(thetap(m-2*i,0)) for i in range(floor(m/2)+1)])\n",
    "\n",
    "# According to Theorem 1.1, Hb_{s_2*s_1*\\theta'(m,0)} = Hb_{s_2}Hb_{s_1}Hb_{\\theta'(m,0)} -Hb_{\\theta'(m,0)} \n",
    "# We would sometimes call s_1=s, s_2=t, s_0=u during the course of our work\n",
    "def Hbtsthetap(m):\n",
    "    return Hb(2)*Hb(1)*HbThetap(m)-HbThetap(m)\n",
    "\n",
    "# premultiply \\theta(m,0) by s_1s_2s_0\n",
    "def HbstuTheta(m):\n",
    "    return((Hb(1)*Hb(2)*Hb(0)-Hb(0)-v^(-1)-v)*Hbtheta(m))\n",
    "\n",
    "# now with t_m multiplied on the right (the element that makes \\theta(m,0) grow)\n",
    "def HbstuThetat_m(m):\n",
    "    if m%2==0:\n",
    "        return(HbstuTheta(m)*Hb(0))\n",
    "    else:\n",
    "        return(HbstuTheta(m)*Hb(1))\n",
    "\n",
    "#similar ideas  \n",
    "\n",
    "def HbsThetat_m(m): # for \\theta(m,0)t_m\n",
    "    return(Hb(1)*HbThetap(m)*Hb((1+m)%2))\n",
    "\n",
    "def HbThetat_ms2(m): # for \\theta(m,0)t_m s_2\n",
    "    if m%2==0:\n",
    "        return(HbTheta(m)*Hb(0)*Hb(2)-HbTheta(m))\n",
    "    else:\n",
    "        return(HbTheta(m)*Hb(1)*Hb(2)-HbTheta(m))\n",
    "\n",
    "def Hbs0theta(m): # for s_0\\theta(m,0)\n",
    "    return(Hb(0)*HbTheta(m))\n",
    "\n",
    "#Supp(m,n) as defined in the introduction of the paper\n",
    "def S(m,n):\n",
    "    if m%2!=0:\n",
    "        return {(m-2*i,n-j) for i in range(floor(m/2)+1) for j in range(n+1)}\n",
    "    else:\n",
    "        return {(m-2*i,n-j) for i in range(floor(m/2)) for j in range(n+1)}.union({(0,n-2*i) for i in range(floor(n/2)+1)})\n",
    "\n",
    "#Direct description of canonical basis elts for thetas as per Theorem 1.1\n",
    "def Hbtheta(m,n):\n",
    "    return sum([v^((m-a)+2*(n-b))*N(theta(a,b)) for (a,b) in S(m,n)])\n",
    "\n",
    "# Now we go to the rest of Theorem 1.1. l and r will represent the x and y terms in that theorem\n",
    "def Hbamigo(m,n,l,r): # 'amigos de theta' are the elements corresponding to alcoves that are either neighbors in the fundamental chamber, \n",
    "# or the corresponding alcoves in the other chambers of the big region of these\n",
    "    if l==1: #multiply on left by identity\n",
    "        if r==1: # on right by identity\n",
    "            return Hbtheta(m,n) # no change\n",
    "        elif r==s0 or r==s1: # right multiply by t_m, etc.\n",
    "            return Hbtheta(m,n)*C[r]\n",
    "        elif r==s0*s2 or r==s1*s2:\n",
    "            return Hbtheta(m,n)*C[r]-Hbtheta(m,n)\n",
    "        elif r==s0*s2*s1:\n",
    "            return Hbamigo(m,n,1,s0*s2)*C[1]-Hbtheta(m,n)*C[0]\n",
    "        elif r==s1*s2*s0:\n",
    "            return Hbamigo(m,n,1,s1*s2)*C[0]-Hbtheta(m,n)*C[1]\n",
    "        else:\n",
    "            print('Bad input, try again.')\n",
    "    elif l==s0:\n",
    "        return Hb(0)*Hbamigo(m,n,1,r)\n",
    "    elif l==s2*s0:\n",
    "        return Hb(2)*Hbamigo(m,n,s0,r)-Hbamigo(m,n,1,r)\n",
    "    elif l==s1*s2*s0:\n",
    "        return Hb(1)*Hbamigo(m,n,s2*s0,r)-Hbamigo(m,n,s0,r)\n",
    "    else:\n",
    "        print('Bad input, try again.')\n",
    "\n",
    "# This tool can help us see the explicit form of the friends of thetas\n",
    "# For example:\n",
    "for m in range(6):\n",
    "    rwN(Hbamigo(m,0,s0,1))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Careful! Already at length 16 this may take quite a long time.\n",
    "T(C[2,1,2,0,2,1,2,0,2,1,2,0,2,1,2,0])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "KL->Standard completed in 1.1920928955078125e-06 seconds.\n",
      "Success! Rewrite took 0.20458483695983887 seconds.\n",
      "1*N[2, 1, 0, 2] + v*N[2]\n"
     ]
    }
   ],
   "source": [
    "# The costly step is not the conversion of the canonical basis element, but its computation\n",
    "rwN(T(C[2,1,0,2]))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "SageMath 9.5",
   "language": "sage",
   "name": "sagemath"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.10.6"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}