class EHMM:
    # Initializes the class for an EHMM data.  Constructs functions to return the various pieces of data
    def __init__(self, alpha, beta, r=None, s=None, t=None, l=1):
        # Check that the inputs are of the correct type/
        if not isinstance(alpha, list):
            raise TypeError("First argument should be a list of parameters which appear in the numerator of the hypergeometric coefficients")
        if not isinstance(beta, list):
            raise TypeError("Second argument should be a list of parameters which appear in the numerator of the hypergeometric coefficients")
        for a in alpha:
            if not (str(a.parent())=='Integer Ring' or str(a.parent())=='Rational Field'):
                raise TypeError("Hypergeometric parameters must be rational numbers")
        for b in beta:
            if not (str(b.parent())=='Integer Ring' or str(b.parent())=='Rational Field'):
                raise TypeError("Hypergeometric parameters must be rational numbers")
        # If beta has one fewer, assume that the additional 1 arising from k! is not being included and add 1
        # For example, [1/2,1/2],[1] could be given as input, this will the datum to [1/2,1/2], [1,1]
        if len(beta)+1==len(alpha):
            beta.append(QQ(1))
        # If the parameter r is given, first checks that it is rational and then adds to alpha,
        # the class will still know what the distinguished parameter r was
        if r!=None:
            if not (str(r.parent())=='Integer Ring' or str(r.parent())=='Rational Field'):
                raise TypeError("Hypergeometric parameters must be rational numbers")
            alpha.append(r)
        # If the parameter r is given, first checks that it is rational and then adds to alpha,
        # the class will still know what the distinguished parameter r was
        if s!=None:
            if not (str(s.parent())=='Integer Ring' or str(s.parent())=='Rational Field'):
                raise TypeError("Hypergeometric parameters must be rational numbers")
            beta.append(s)
        # If alpha and beta don't have the same length after ammending as above, throw an error
        if len(beta)!=len(alpha):
            raise ValueError("Can only consider generalized hypergeometric functions of the form nFn-1")
        self.alpha=alpha
        self.beta=beta
        self.r=r
        self.s=s
        self.t=t
        self.l=l
        
   # Print Hauptmodul
    def tq(self):
        A1=self.alpha;B1=self.beta;n=len(A1)
        A=[A1[i] for i in srange(len(A1)-1)];B=[B1[i] for i in srange(len(B1)-1)]
        r=A1[n-1]; s=B1[n-1]
        if self.t==None:
            for data in HauptTable:
                if [A,B]==data[0]:
                    t=R(data[1])
        return(t) 
        
    # Compute the demoninator M of a EHMM object
    def denominator(self):
        M=1
        for a in self.alpha:
            M=lcm(M,a.denominator())
        for b in self.beta:
            M=lcm(M,b.denominator())
        return(M)
    
    # Returns the value of gamma(HD)
    def gamma(self):
        S=-1
        # We're assuming that alpha and beta have the same length in the construction of the EHMM object, so we only need one for loop
        for i in range(len(self.alpha)):
            S=S+self.beta[i]-self.alpha[i]
        return(S)

    # Compute q-expansion
    def q_expansion(self):
        q=var('q')
        A1=self.alpha;B1=self.beta;n=len(A1)
        A=[A1[i] for i in srange(len(A1)-1)];B=[B1[i] for i in srange(len(B1)-1)]
        r=A1[n-1]; s=B1[n-1]
        if self.t==None:
            for data in HauptTable:
                if [A,B]==data[0]:
                    t=R(data[1])
        if t==None:
            print('Hypergeometric datum does not have corresponding Hauptmodul in Table 1')
            return()        
        deg=0
        while ((t)[deg]==0):       #Locating first leading coefficient exponent
          deg=deg+1  
        C=(t)[deg]                 #leading coefficient of t
        M=denominator(deg*(r-1)+1) # M=denominator(deg*r)
        n=r*M   # n=numerator(deg*r)
        rmf=(((t/C/R(q**deg))**(r-1)*(1-t)**(s-r-1)*Hyp(A,B,t,N)*diff(t,q)).O(50*M)/C/deg) 
        if M==1:
            return R(rmf*q**(deg*(r-1)+1))
        else:    
            tar=sum(rmf[i]*q**(M*i+n) for i in srange(50*M))
            return R(tar)
    def q_expansion(self):
        A1=self.alpha;B1=self.beta
        return q_exp(A1,B1)
    
#### Compute conjugates of (r,s)
    def conjugate(self):
        A=self.alpha; B=self.beta
        Conj=[]
        M=self.denominator()
        for i  in srange(1,M):
         if gcd(i,M)==1:
            Conj.append([[fr(i*a) for a in A],[fr0(i*a) for a in B]])
        return Conj
        
#### Expansions for all conjugates
    def q_conjugate_exp(self):        
        CC=self.conjugate()
        return table([q_exp(c[0],c[1]).O(20) for c in CC])
        
#### zigzag diagram for all conjugates        
    def conjugate_zigzag(self):#conjugate_zz(self):
        CC=self.conjugate();        n=len(CC);
        zz=[zigzag(c[0],c[1]) for c in CC]
        if n==1:
            return (zz[0]).show()
        else:    
            List=[]
            for i in srange(ZZ(floor(n/2))):
                List.append([zz[i*2],zz[i*2+1]])
            return graphics_array(List).show()
        
#### normalized character sums
    def normalized_Hp(self):
            M=self.denominator()
            List=[['p','Normalized Hp']];A1=self.alpha;B1=self.beta;
            n=len(A1);r=A1[n-1];
            A=[A1[i] for i in srange(len(A1)-1)];B=[B1[i] for i in srange(len(B1)-1)]
            if self.t==None:
             for data in HauptTable:
                if [A,B]==data[0]:
                    t=R(data[1])
            if t==None:
             print('Hypergeometric datum does not have corresponding Hauptmodul in Table 1')
             return()        
            deg=0
            while ((t)[deg]==0):       #Locating first leading coefficient exponent
                deg=deg+1  
            C1=(t)[deg]                 #leading coefficient of t
            for p in primes(15*M):  #can adjust 15 to compute more MF's coefficients
              if p%M==1:
                 chi=DirichletGroup(p)[(p-1)*r]
                 ss=(-1)**(n-1)*PP(A1,B1,1,p)*chi(C1)
                 List.append([p,ss])
            return table(List)

    def coeff_p_order(self):
        A=self.alpha
        B=self.beta
        Z=Qp(p)
        N1=50
        List=[];Plot=[]
        for k in srange(N1):
            List.append(Z(AHD(k,A,B)).valuation())
            
        for k in srange(N1-1):
            Plot.append(polygon([(k,0),(k,List[k]),(k+1,List[k]),(k+1,0)],color='grey'))
        return sum(p for p in Plot)+text("ord_p of hyper. coeffs for datum",(12,-3))+text([A,B],(30,-3))+text(" for p=",(40,-3))+text(p,(43,-3))