import numpy as np
from sympy import Symbol
from sympy.solvers import solve
from scipy.optimize import fsolve
from functools import lru_cache


class ARMA(object):
    def __init__(self, process, m, n):
        self.process = process
        self.m = m
        self.n = n
        self.z = None

        self.betas = np.array(list(fsolve(self.beta_equations_system, np.ones(self.m), factor=0.1)))
        result = fsolve(self.alpha_equations_system, np.ones(self.n * 2 + 2), factor=0.1, full_output=True)
        self.alphas = np.array(list(result[0][0:self.n + 1]))
        self.norm = np.linalg.norm(result[1]['fvec'])

        self.is_model_exists = not (result[2] != 1 and self.norm >= 1e-1)

    @lru_cache(maxsize=None)
    def normalized_correlation(self, k):
        k = abs(k)
        if k <= self.m + self.n:
            return self.process.normalized_correlation(k)
        else:
            ncf = 0
            for i in range(1, self.m + 1):
                ncf += self.betas[i - 1] * self.normalized_correlation(k - i)
            return ncf

    @lru_cache(maxsize=None)
    def std(self, count=10):
        result = 0
        for i in range(1, count + 1):
            result += (self.normalized_correlation(i) - self.process.normalized_correlation(i)) ** 2
        return result

    @lru_cache(maxsize=None)
    def is_model_stable(self):
        if self.z is None:
            z = Symbol("z")
            expression = 1
            for i in range(self.m):
                expression -= self.betas[i] * (z ** -(i + 1))
            self.z = solve(expression, z)
        if self.m == 2 and self.n == 3:
            print(list(self.z))
        return all(abs(z) < 1 for z in self.z)

    def beta_equations_system(self, beta):
        result = []
        for i in range(self.n + 1, self.n + self.m + 1):
            current = -self.process.correlation_function(i)
            for j in range(1, self.m + 1):
                current += beta[j - 1] * self.process.correlation_function(i - j)
            result.append(current)
        return result

    def alpha_equations_system(self, factors):
        result = []

        for i in range(0, self.n + 1):
            current = -self.process.correlation_function(i)
            for j in range(1, self.m + 1):
                current += self.betas[j - 1] * self.process.correlation_function(i - j)
            for j in range(i, self.n + 1):
                current += factors[j] * factors[self.n + 1 + j - i]
            result.append(current)

        for i in range(0, self.n + 1):
            current = factors[i] - factors[self.n + 1 + i]
            for j in range(1, min(i, self.m) + 1):
                current += self.betas[j - 1] * factors[self.n + 1 + i - j]
            result.append(current)

        return result