# -*- coding: utf-8 -*-
"""
Spyder Editor

This is a temporary script file.
"""
import glob
import os
import sympy
import re

path = 'C:\Users\Gebruiker\Documents\Radboud\Informatica\Jaar 2\Periode 1\Combinatoriek\ComputerAssignment'

import numpy as np
from numpy.polynomial import Polynomial as P
from numpy.linalg import solve, lstsq

def get_roots(poly):
    r = poly.r
    r = r[np.isreal(r)]
    r = [np.real(x) for x in r]
    m = [1 for i in range(0,len(r))]
    # calculate multiplicity using repeated differentiating
    # the derivative has multiplicity n-1
    while poly.o > 0:
        poly = np.polyder(poly)
        tmp = poly.r
        tmp = list(set(tmp[np.isreal(tmp)])) # use set to remove duplicates to avoid double counting
        for i in range(0,len(r)):
            for x in tmp:
                if(np.isclose(x,r[i])):
                    m[i]+=1
    # remove duplicates in array, sometimes numpy puts roots multiple times
    if(len(r)==0):
        return r,m
    for x in range(0,len(r)-1):
        for y in range(x+1,len(r)):
            if(np.isclose(r[x],r[y])):
                m[x]+=1
                r[y]='x'
                m[y]='x'
    if 'x' in r:
        r.remove('x')
    if 'x' in m:
        m.remove('x')
    return r,m

def get_constants(r,m,initcond):
    rows = len(initcond)
    cols = np.sum(m)
    mat = np.empty([rows,cols])
    for n in range(0,len(initcond)):
        i = 0
        for k in range(0,len(r)):
            for l in range(0,m[k]):
                mat[n][i]=(n**l)*(r[k]**n)
                i+=1
    constants = lstsq(mat,initcond)[0]    # lstsq instead of solve, which makes matrix linearly independent first
    return constants

def get_formula(r,m,constants):
    formula = ""
    i = 0
    for k in range(0,len(r)):
        formula+="("
        for l in range(0,m[k]):
            formula+=str(constants[i])
            if(l==1):
                +"*n"
            elif(l!=0):
                +"*n^"+str(l)
            if(l<m[k]-1):
                formula+="+"
            i+=1
        formula+=")*(" + str(r[k]) + ")^n"
        if(k<len(r)-1):
            formula+="+"
    return formula

def get_max_degree(receq):
    deg = 0
    for (x,y) in receq:
        if y > deg:
            deg = y
    return deg

#klopt niet:
def get_characteristic(receq):
    deg = get_max_degree(receq)
    coef = []
    coef.append(1)
    for i in range(1,deg+1):
        a = 0
        for (x,y) in receq:
            if(y==i):
                a=x*(-1)
        coef.append(a)
    c = np.poly1d(coef)
    print(c)
    return c

def solve_homogenous(receq,initcond,debug=False):
    poly = get_characteristic(receq)
    r,m = get_roots(poly)
    constants = get_constants(r,m,initcond)
    if(debug):
        print "characteristic equation:"
        print poly
        print "roots:"
        print r
        print "multiplicities:"
        print m
        print "constants:"
        print constants
        print ""
    formula = get_formula(r,m,constants)
    return formula

def get_initcond(functions):
    initcond=[]
    for i in range (1, len(functions)):
        function = functions[i]
        ints = map(int, re.findall(r'-?\d*\.{0,1}\d', function)) #finds all integers \d+
        initcond.append(ints[1])
    return initcond

def get_receq(function):
    #if()
    receq = []
    index_s = [m.start() for m in re.finditer('s', function)]
    split_index = 0
    for i in range(0, len(index_s)):
        if(i==0):
            split_index = function.find(")",index_s[i])
            function_part = function[:split_index+1]
        elif(i==len(index_s)-1):
            function_part = function[split_index+1:]
        else:
            split_help = split_index+1
            split_index = function.find(")", split_index+1)
            function_part = function[split_help:split_index+1]
        ints = map(int, re.findall(r'-?\d*\.{0,1}\d', function_part))
        if (len(ints) == 1):
            if(function_part.find('s')==0):
                receq.append((1,ints[0]*-1))
            elif(function_part[function_part.find('s')-1] == '+'):
                receq.append((1,ints[0]*-1))
            elif(function_part[function_part.find('s')-1] == '-'):
                receq.append((-1,ints[0]*-1))
        else:
            receq.append((ints[0],ints[1]*-1))
    print(receq)
    return receq

#open file and split data
for filename in glob.glob(os.path.join(path, 'comass??.txt')):
    homogenous = False
    exponential = False
    polynomial = False
    with open(filename,"r") as file:
        data = file.read()              #put file data in data variable
    data = data[data.index("[")+1:]     #remove the eqs[=] part of the string
    data = data[:data.index("]")]
    data = data.replace(" ","")
    functions = data.split(",")         #put the seperate equations in an array
    functions[0] = functions[0][functions[0].index("=")+1:] #remove the s[n]= part from the rec eq
    hom = raw_input(functions[0] + str(' Is it homogenous, y or n?\n'))
    if(hom=="y"):
        homogenous = True
    else:
        exp = raw_input('Is it exponential, y or n?\n')
        if(exp=="y"):
            exponential = True
        poly = raw_input('Is it polynomial, y or n?\n')
        if(poly=="y"):
            polynomial = True
    if(homogenous):
        #initcond = [1,1] # [n0, n1, ... , nk]
        #receq = [(1,2),(1,1)] # (constant,degree)
        initcond=get_initcond(functions)
        receq = get_receq(functions[0])
        formula = solve_homogenous(receq,initcond,True)

    print "formula:"
    print formula

    #write the solution the right file:
    filename = filename[:-4] + str("-dir.txt")
    file = open(filename, "w")
    file.write("sdir := n -> " + str(formula))
    file.close()