import numpy as np 
import math
from typing import List
import os
import matplotlib.pyplot as plt

class Interp1d:
    def __init__(self, xs : List[float], ys : List[float]):
        xs, ys = np.array(xs), np.array(ys)
        # ascending order
        inds = np.argsort(xs)
        self.xs, self.ys = xs[inds], ys[inds]
        self.len = len(self.xs)

    def __call__(self, x):
        lowerboundp = 0
        # optain the lowerbound
        for i, xi in enumerate(self.xs):
            if x >= xi:
                lowerboundp = i
            else:
                break
        if lowerboundp < self.len - 1:
            upperboundp = lowerboundp + 1
        else:
            lowerboundp, upperboundp = self.len - 2, self.len - 1

        x0, y0, x1, y1 = self.xs[lowerboundp], self.ys[lowerboundp], self.xs[upperboundp], self.ys[upperboundp]
        if x1 == x0:
            return (y0 + y1) / 2.0

        return y0 + (x-x0)/(x1-x0)*(y1-y0)

        
        

def linspace(start : float, end : float, n : int):
    samples = []
    step = (end - start) / (n-1)
    for i in range(n):
        samples.append(start + float(i)*step)
    return samples

def zeros(n : int):
    return [0] * n

def aczwave(amplitude : float, length : int, 
    carrierFreq : float, carrierPhase : float, dragAlpha : float,
    thf : float, thi : float, lam2 : float, lam3 : float):

    t = linspace(0, 1, length)
    han2 = []
    for k, x in enumerate(t):
        han2.append(
            (1-lam3)*(1-math.cos(2.0*math.pi*x)) +
            lam2*(1-math.cos(4*math.pi*x)) +
            lam3*(1-math.cos(6*math.pi*x))
        )
    maxHan2 = max(han2)

    ths1 = []
    for k in range(length):
        ths1.append(
            thi + (thf-thi)*han2[k]/maxHan2
        )
    t1u = zeros(length)
    for k, v in enumerate(t1u):
        if k < (length - 1):
            t1u[k+1] = v + math.sin(ths1[k])/float(length-1)

    for k, v in enumerate(t):
        t[k] = v * t1u[length-1]

    th = Interp1d(t1u, ths1)
    th0 = 1.0 / math.tan(th(t[0]))

    thval = []
    for k in range(length):
        thval.append(
            1.0/math.tan(th(t[k])) - th0 
        )
    thmin = min(thval)

    samples = []
    for k in range(length):
        env = thval[k] * amplitude / thmin 
        samples.append(complex(env, 0))

    return samples


def test():
    amplitude = 26214
    length = 30
    carrierFreq = 0
    carrierPhase = 0.000000
    dragAlpha = 0.000000
    thf = 0.864
    thi = 0.05
    lam2 = -0.18
    lam3 = 0.04

    data = aczwave(
        amplitude, length, carrierFreq,
        carrierPhase, dragAlpha,
        thf, thi, lam2, lam3,
    )
    for c in data:
        print(c.real, c.imag)
    return data


class Benchmark:
    def __init__(self, num_samplings : int):
        self.data_dir = "data"
        self.num_samplings = num_samplings
        self.params_dict = {}
        self.gt_dict = {} 
        self.load_params()
        self.load_gt()

    def load_params(self):
        for i in range(self.num_samplings):
            file = os.path.join(self.data_dir, "aczgo_param_{}.log".format(i))
            with open(file, 'r') as f:
                lines = f.readlines()
            params = {}
            for line in lines:
                key, value = line.split(", ")
                if key=="length":
                    value = int(value)
                else:
                    value = float(value)
                params[key] = value

            self.params_dict[i] = params

    def eval(self, idx):
        params = self.params_dict[idx]

        amplitude = params["amplitude"]
        length = params["length"]
        carrierFreq = params["carrierFreq"]
        carrierPhase = params["carrierPhase"]
        dragAlpha = params["dragAlpha"]
        thf = params["thf"] 
        thi = params["thi"] 
        lam2 = params["lam2"]
        lam3 = params["lam3"]

        data = aczwave(
            amplitude, length, carrierFreq,
            carrierPhase, dragAlpha,
            thf, thi, lam2, lam3,
        )
        xs, ys = [], []
        for c in data:
            xs.append(c.real)
            ys.append(c.imag)

        return (xs, ys)

    def load_gt(self):
        for i in range(self.num_samplings):
            file = os.path.join(self.data_dir, "aczgo_result_{}.log".format(i))
            xs, ys = [], []
            with open(file, 'r') as f:
                lines = f.readlines()
            for line in lines:
                x, y = line.split(", ")
                x, y = float(x), float(y)
                xs.append(x)
                ys.append(y)
            self.gt_dict[i] = (xs, ys)

    def test(self, idx):
        def check_valid(vs):
            return all(map(lambda x:not np.isnan(x) and not np.isinf(x), vs))
        def max_ab_dis(xs, bxs):
            return np.abs(np.array(bxs) - np.array(xs)).max()


        (bxs, bys) = self.gt_dict[idx]
        if check_valid(bxs) and check_valid(bys):
            xs, ys = self.eval(idx)
            return (max_ab_dis(xs, bxs), max_ab_dis(ys, bys))
        else:
            return "not valid"

    def test_all(self):
        for i in range(self.num_samplings):
            print(self.test(i))


if __name__ == "__main__":
    b = Benchmark(11)
    print(b.test_all())



# data = test()
#
# np.savetxt('D:/Work/TailCorr/acz_750.csv', data, delimiter=' ')
# plt.figure()
# plt.plot(data)
# plt.show()