import random
import matplotlib.pyplot as plt
import os
import numpy as np


class Polynomial:
    def __init__(self, degree, noise_type, noise_std, param_max=2.5):
        self.params = np.random.uniform(
            low=-param_max, high=param_max, size=degree+1)
        self.degree = degree
        self.noise_std = noise_std
        self.noise_type = noise_type

    def predict_no_noise(self, x):
        pred = np.zeros(x.shape)
        for i in range(self.degree+1):
            pred += x**(i) * self.params[i]
        return pred
    
    def predict_with_noise(self,x, scale):
        pred = np.zeros(x.shape)
        for i in range(self.degree+1):
            pred += x**(i) * self.params[i]
        if self.noise_type == "uniform":
            noise = (np.random.uniform(
                size=x.shape, high=self.noise_std) - self.noise_std/2)*scale
        else:
            noise = np.random.normal(size=x.shape, scale=self.noise_std)*scale
        return pred + noise

# univariate continuous spline sampling
# uniformly we sample the number of bins as [min_segments,max_segments] 
# each bin contains [min_points, max_points] datapoints
# point_sampling specifies either "uniform" or exponential sampling of number of samples per bin
# noise type may be either gaussian or uniform
# the noise  variance is uniformly drawn between [noise_low,noise_high]
# heteroscedasticity is introducted through the step parameters, where the variance between
# each bin is prev_var +- [step_low, step_high]. Set both steps to 0 for homoscedastic
# if continuous is true C0 continuity is enforced
# polynomial_degree specifies the degree of bin functions which is sampled
 
def make_the_data(min_segments, max_segments, min_points, max_points, point_sampling, noise_type, noise_low, noise_high, continuous,step_low,step_high, polynomial_degree=3):
    num_breakpoints = np.random.randint(min_segments, max_segments+1)
    if point_sampling == "uniform":
        num_points = [np.random.randint(min_points, max_points+1)
                  for i in range(num_breakpoints)]
    else:
        num_points = [round(2**np.random.uniform(min_points, max_points))
                    for i in range(num_breakpoints)]
    initial_var = np.random.uniform(noise_low,noise_high)
    
    variance_staircase = []
    for i in range(num_breakpoints):
        step = np.random.uniform(step_low,step_high)
        sign = -1 if np.random.uniform() < 0.5 else 1
        if initial_var < step:
            sign = 1
        initial_var += step*sign
        variance_staircase.append(initial_var)
    noise_type = random.choice(["uniform","gaussian"])
    
    segment_functions = [Polynomial(np.random.randint(
            1, polynomial_degree+1), noise_type, variance_staircase[i]) for i in range(num_breakpoints)]
        
    bps = np.random.uniform(low=1, high=5, size=num_breakpoints)
    generatedDataX = []
    generatedDataY = []
    start = 0
    sample_start = 0
    sample_end = 2
    min_y = None
    max_y = None
    # figure out domain to get percentage noise
    for i in range(num_breakpoints):
        segment_fun = segment_functions[i]
        end = start + bps[i]
        segment_sample = np.random.uniform(
            low=sample_start, high=1, size=num_points[i])
        segment_y = segment_fun.predict_no_noise(2*segment_sample)
        if min_y is None or np.amin(segment_y) < min_y:
            min_y = np.amin(segment_y)
        if max_y is None or np.amax(segment_y) > max_y:
            max_y = np.amax(segment_y)
    scale = max_y-min_y

    for i in range(num_breakpoints):
        segment_fun = segment_functions[i]
        if start > 0:
            if continuous:
                now = segment_fun.predict_no_noise(np.array(sample_start))
            else:
                now = segment_fun.predict_with_noise(np.array(sample_start),scale)
            segment_fun.params[0] = segment_fun.params[0] - now + oldIntercept
        end = start + bps[i]

        segment_sample = np.random.uniform(
            low=sample_start, high=1, size=num_points[i])
        segment_x = segment_sample * (end-start) + start
        segment_y = segment_fun.predict_with_noise(2*segment_sample,scale)

        generatedDataX.append(segment_x)
        generatedDataY.append(segment_y)
        start = end
        oldIntercept = segment_fun.predict_no_noise(np.array(sample_end))
    return generatedDataX, generatedDataY, bps

#ITCI 21 data



points = (25, 25, "uniform", "few_points")
noise = (0.2, 0.3, "uniform", "low_gaussian_noise")

for low,high in [(0,0),(0.01,0.04),(0.04,0.08),(0.08,0.16),(0.16,0.32)]:
    pairmeta = []
    for i in range(1, 101):
        X, Y, _ = make_the_data(3, 3, points[0],
                                points[1], points[2], noise[2], noise[0], noise[1], "Spline",low,high)
        X = np.concatenate(X)
        Y = np.concatenate(Y)
        direction = np.random.randint(2)
        if direction == 1:
            dataset = np.stack((X, Y)).T
        else:
            dataset = np.stack((Y, X)).T
        pairmeta.append(direction)
        if low == 0:
            newfolder = "GeneratedData/HeteroscedasticPairs/Homoscedastic/"
        else:
            newfolder = "GeneratedData/HeteroscedasticPairs/Step_{:d}_{:d}/".format(round(low*100),round(high*100))
        if not os.path.exists(newfolder):
            os.makedirs(newfolder)
        np.savetxt(newfolder + "/pair_" +
                    str(i)+".txt", dataset)
        if i < 2:
            plt.figure()
            plt.scatter(X, Y)
    np.savetxt(newfolder+"pairs_gt.txt",pairmeta)
exit()