from penguins.Model import Model as model
import numpy as np
import matplotlib.pyplot as plt
import scipy.stats as stats
import pandas as pd

#region Creating Randomized Complicated Data
def funcComplicated(x, A, B, C):
    return ((A-x)**3)*np.sin(B*x)/(1-np.exp(C*x))
vecComplicated = np.vectorize(funcComplicated)

def randomize(func,range, **kwargs):
    data = pd.DataFrame()
    data['X'] = np.arange(range[0],range[1], .01)

    size = len(data)
    offset = 0.05
    params = {}
    for key in kwargs:
        params[key] = stats.norm.rvs(loc=kwargs[key], scale=kwargs[key]*offset,
                   size=size)

    vecFunc = np.vectorize(func)

    data['Y'] = vecFunc(data['X'],*params.values())
    return data
#endregion

def test_SetParams_CaculatedParams():
    """
    Here I'm simply testing the Quiet and Verbose modes. This tests the fit 'mode' setting
    """

    init = {'A': 0, 'B': 2.50, 'C': .70}
    data = randomize(funcComplicated, [0.01, 2 * np.pi], A=1, B=3, C=.25)
    calculatedParams = {
        'A'      : lambda args: args['A'],
        'B'      : lambda args: args['B'],
        'C'      : lambda args: args['C'],
        'A*B'    : lambda args: args['A']*args['B'],
        'A+B'    : lambda args: args['A']+args['B'],
        'A^C': lambda args: args['A']**args['C']
    }
    CompMod = model(funcComplicated)
    CompMod.setParams(init,calculatedParams=calculatedParams)
    CompMod.fit(data['X'], data['Y'])
    CompMod.print(data['X'], data['Y'])

def test_Fit_ReturnParams():
    """
    Testing all Algorithms and Methods
    """
    init = {'A': 0, 'B': 2.50, 'C': .70}
    data = randomize(funcComplicated,[0.01,2*np.pi], A = 1, B = 3, C = .25)
    CompMod = model(funcComplicated)

    CompMod.setParams(init)
    CompMod.returnParameters()
    CompMod.fit(data['X'], data['Y'])
    CompMod.returnParameters()

