import numpy as np
import h5py
def createDataSetOld():
    new_file_path_train = "testtestssss.h5"
    new_file_path_test = "traintrain.h5"
    #for each shpae
    boundaries =[0.5, 0.5,0,5, 1.8]
    epsilons = [0.05, 0.2, 0.2, 0.2]
    with h5py.File(new_file_path_train, "w") as new_hdf5_train_file:
        point_clouds_group = new_hdf5_train_file.create_group("point_clouds")

        # Train data generation for each label
        addDataToSet(point_clouds_group=point_clouds_group,gaussian_curv=0, mean_curv=0, label=0, boundary=1, epsilon=0.05, max_curv=5, counter=0, size=1.5, amount_of_pcl=10000)
        print(f'Finished train flat surfaces')

        addDataToSet(point_clouds_group=point_clouds_group,gaussian_curv=1, mean_curv=1, label=1, boundary=1, epsilon=0.05, max_curv=5, counter=10000,
                     size=3, amount_of_pcl=2500)
        print(f'Finished train parabolic peak surfaces')
        addDataToSet(point_clouds_group=point_clouds_group,gaussian_curv=1, mean_curv=-1, label=1, boundary=1, epsilon=0.05, max_curv=5, counter=12500,
                     size=3, amount_of_pcl=2500)
        print(f'Finished train parabolic pit surfaces')
        addDataToSet(point_clouds_group=point_clouds_group,gaussian_curv=1, mean_curv=-1, label=1, boundary=1, epsilon=0.05, max_curv=5, counter=15000,
                     size=3, amount_of_pcl=2500, angle=1)
        print(f'Finished train parabolic CORNERS')
        addDataToSet(point_clouds_group=point_clouds_group,gaussian_curv=1, mean_curv=-1, label=1, boundary=1, epsilon=0.05, max_curv=3, counter=17500,
                     size=3, amount_of_pcl=2500, radius=1)
        print(f'Finished train parabolic SPHERES')

        addDataToSet(point_clouds_group=point_clouds_group,gaussian_curv=0, mean_curv=1, label=2, boundary=1, epsilon=0.05, max_curv=5, counter=20000,
                     size=2, amount_of_pcl=2500)
        print(f'Finished train ridge surfaces')
        addDataToSet(point_clouds_group=point_clouds_group,gaussian_curv=0, mean_curv=-1, label=2, boundary=1, epsilon=0.05, max_curv=5, counter=22500,
                     size=2, amount_of_pcl=2500)
        print(f'Finished train valley surfaces')
        addDataToSet(point_clouds_group=point_clouds_group,gaussian_curv=0, mean_curv=1, label=2, boundary=1, epsilon=0.05, max_curv=5, counter=25000,
                     size=2, amount_of_pcl=2500, angle=1)
        print(f'Finished train ridge ANGLES')
        addDataToSet(point_clouds_group=point_clouds_group,gaussian_curv=0, mean_curv=-1, label=2, boundary=1, epsilon=0.05, max_curv=3, counter=27500,
                     size=2, amount_of_pcl=2500, radius=1)
        print(f'Finished train valley Cylinders')

        # addDataToSet(point_clouds_group=point_clouds_group,gaussian_curv=-1, mean_curv=-33, label=3, boundary=4, epsilon=0.05, max_curv=5, counter=30000,
        #              size=2.5, amount_of_pcl=10000)
        # print(f'Finished train saddle surfaces')

        # addDataToSet(point_clouds_group=point_clouds_group,gaussian_curv=0, mean_curv=0, label=4, boundary=1, epsilon=0.05, max_curv=5, counter=40000,
        #              size=3, amount_of_pcl=2500)
        # print(f'Finished train HALFSPACE flat surfaces')
        # addDataToSet(point_clouds_group=point_clouds_group,gaussian_curv=1, mean_curv=1, label=4, boundary=1, epsilon=0.05, max_curv=5, counter=42500, size=3, amount_of_pcl=625)
        # print(f'Finished train HALFSPACE parabolic peak surfaces')
        # addDataToSet(point_clouds_group=point_clouds_group,gaussian_curv=1, mean_curv=-1, label=4, boundary=1, epsilon=0.05, max_curv=5, counter=43125,
        #              size=3, amount_of_pcl=625)
        # print(f'Finished train HALFSPACE parabolic pit surfaces')
        # addDataToSet(point_clouds_group=point_clouds_group,gaussian_curv=1, mean_curv=1, label=4, boundary=1, epsilon=0.05, max_curv=5, counter=43750, size=3, amount_of_pcl=625,
        #              angle=1, edge=1)
        # print(f'Finished train HALFSPACE parabolic CORNERS')
        # addDataToSet(point_clouds_group=point_clouds_group,gaussian_curv=1, mean_curv=-1, label=4, boundary=1, epsilon=0.05, max_curv=3, counter=44375,
        #              size=3, amount_of_pcl=625, radius=1, edge=1)
        # print(f'Finished train HALFSPACE parabolic SPHERES')
        # addDataToSet(point_clouds_group=point_clouds_group,gaussian_curv=0, mean_curv=1, label=4, boundary=1, epsilon=0.05, max_curv=5, counter=45000, size=3, amount_of_pcl=625)
        # print(f'Finished train HALFSPACE ridge surfaces')
        # addDataToSet(point_clouds_group=point_clouds_group,gaussian_curv=0, mean_curv=-1, label=4, boundary=1, epsilon=0.05, max_curv=5, counter=45625,
        #              size=3, amount_of_pcl=625)
        # print(f'Finished train HALFSPACE valley surfaces')
        # addDataToSet(point_clouds_group=point_clouds_group,gaussian_curv=0, mean_curv=1, label=4, boundary=1, epsilon=0.05, max_curv=5, counter=46250, size=3, amount_of_pcl=625,
        #              angle=1, edge=2)
        # print(f'Finished train HALFSPACE ridge ANGLES')
        # addDataToSet(point_clouds_group=point_clouds_group,gaussian_curv=0, mean_curv=-1, label=4, boundary=1, epsilon=0.05, max_curv=3, counter=46875,
        #              size=3, amount_of_pcl=625, radius=1, edge=2)
        # print(f'Finished train HALFSPACE valley CYLINDERS')
        # addDataToSet(point_clouds_group=point_clouds_group,gaussian_curv=-1, mean_curv=-33, label=4, boundary=4, epsilon=0.05, max_curv=5, counter=47500,
        #              size=3, amount_of_pcl=2500)
        # print(f'Finished train HALFSPACE saddle surfaces')

    with h5py.File(new_file_path_test, "w") as new_hdf5_test_file:
        point_clouds_group = new_hdf5_test_file.create_group("point_clouds")
        # Test data generation for each label
        addDataToSet(point_clouds_group=point_clouds_group,gaussian_curv=0, mean_curv=0, label=0, boundary=1, epsilon=0.05, max_curv=5, counter=0, size=1.5, amount_of_pcl=1000)
        print(f'Finished test flat surfaces')

        addDataToSet(point_clouds_group=point_clouds_group,gaussian_curv=1, mean_curv=1, label=1, boundary=1, epsilon=0.05, max_curv=5, counter=1000, size=3, amount_of_pcl=250)
        print(f'Finished test parabolic peak surfaces')
        addDataToSet(point_clouds_group=point_clouds_group,gaussian_curv=1, mean_curv=-1, label=1, boundary=1, epsilon=0.05, max_curv=5, counter=1250, size=3, amount_of_pcl=250)
        print(f'Finished test parabolic pit surfaces')
        addDataToSet(point_clouds_group=point_clouds_group,gaussian_curv=1, mean_curv=-1, label=1, boundary=1, epsilon=0.05, max_curv=5, counter=1500, size=3, amount_of_pcl=250,
                     angle=1)
        print(f'Finished test parabolic CORNERS')
        addDataToSet(point_clouds_group=point_clouds_group,gaussian_curv=1, mean_curv=-1, label=1, boundary=1, epsilon=0.05, max_curv=3, counter=1750, size=3, amount_of_pcl=250,
                     radius=1)
        print(f'Finished test parabolic SPHERES')

        addDataToSet(point_clouds_group=point_clouds_group,gaussian_curv=0, mean_curv=1, label=2, boundary=1, epsilon=0.05, max_curv=5, counter=2000, size=2, amount_of_pcl=250)
        print(f'Finished test ridge surfaces')
        addDataToSet(point_clouds_group=point_clouds_group,gaussian_curv=0, mean_curv=-1, label=2, boundary=1, epsilon=0.05, max_curv=5, counter=2250, size=2, amount_of_pcl=250)
        print(f'Finished test valley surfaces')
        addDataToSet(point_clouds_group=point_clouds_group,gaussian_curv=0, mean_curv=1, label=2, boundary=1, epsilon=0.05, max_curv=5, counter=2500, size=2, amount_of_pcl=250,
                     angle=1)
        print(f'Finished test ridge ANGLES')
        addDataToSet(point_clouds_group=point_clouds_group,gaussian_curv=0, mean_curv=-1, label=2, boundary=1, epsilon=0.05, max_curv=3, counter=2750, size=2, amount_of_pcl=250,
                     radius=1)
        print(f'Finished test valley Cylinders')

        # addDataToSet(point_clouds_group=point_clouds_group,gaussian_curv=-1, mean_curv=-33, label=3, boundary=4, epsilon=0.05, max_curv=5, counter=3000,
        #              size=2.5, amount_of_pcl=1000)
        # print(f'Finished test saddle surfaces')

        # addDataToSet(point_clouds_group=point_clouds_group,gaussian_curv=0, mean_curv=0, label=4, boundary=1, epsilon=0.05, max_curv=5, counter=4000, size=3, amount_of_pcl=250)
        # print(f'Finished test HALFSPACE flat surfaces')
        # addDataToSet(point_clouds_group=point_clouds_group,gaussian_curv=1, mean_curv=1, label=4, boundary=1, epsilon=0.05, max_curv=5, counter=4250, size=3, amount_of_pcl=63)
        # print(f'Finished test HALFSPACE parabolic peak surfaces')
        # addDataToSet(point_clouds_group=point_clouds_group,gaussian_curv=1, mean_curv=-1, label=4, boundary=1, epsilon=0.05, max_curv=5, counter=4313, size=3, amount_of_pcl=63)
        # print(f'Finished test HALFSPACE parabolic pit surfaces')
        # addDataToSet(point_clouds_group=point_clouds_group,gaussian_curv=1, mean_curv=1, label=4, boundary=1, epsilon=0.05, max_curv=5, counter=4376, size=3, amount_of_pcl=63,
        #              angle=1, edge=1)
        # print(f'Finished test HALFSPACE parabolic CORNERS')
        # addDataToSet(point_clouds_group=point_clouds_group,gaussian_curv=1, mean_curv=-1, label=4, boundary=1, epsilon=0.05, max_curv=3, counter=4439, size=3, amount_of_pcl=63,
        #              radius=1, edge=1)
        # print(f'Finished test HALFSPACE parabolic SPHERES')
        # addDataToSet(point_clouds_group=point_clouds_group,gaussian_curv=0, mean_curv=1, label=4, boundary=1, epsilon=0.05, max_curv=5, counter=4502, size=3, amount_of_pcl=63)
        # print(f'Finished test HALFSPACE ridge surfaces')
        # addDataToSet(point_clouds_group=point_clouds_group,gaussian_curv=0, mean_curv=-1, label=4, boundary=1, epsilon=0.05, max_curv=5, counter=4565, size=3, amount_of_pcl=63)
        # print(f'Finished test HALFSPACE valley surfaces')
        # addDataToSet(point_clouds_group=point_clouds_group,gaussian_curv=0, mean_curv=1, label=4, boundary=1, epsilon=0.05, max_curv=5, counter=4628, size=3, amount_of_pcl=63,
        #              angle=1, edge=2)
        # print(f'Finished test HALFSPACE ridge ANGLES')
        # addDataToSet(point_clouds_group=point_clouds_group,gaussian_curv=0, mean_curv=-1, label=4, boundary=1, epsilon=0.05, max_curv=3, counter=4691, size=3, amount_of_pcl=63,
        #              radius=1, edge=2)
        # print(f'Finished test HALFSPACE valley CYLINDERS')
        # addDataToSet(point_clouds_group=point_clouds_group,gaussian_curv=-1, mean_curv=-33, label=4, boundary=4, epsilon=0.05, max_curv=5, counter=4754,
        #              size=3, amount_of_pcl=250)
        # print(f'Finished test HALFSPACE saddle surfaces')

def addDataToSet(point_clouds_group, gaussian_curv, mean_curv, label, counter, amount_of_pcl, boundary,max_curv, epsilon,size, angle=0, radius=0, edge=0, update=False):
    default_max_angle = 120
    default_min_angle = 60
    constant = max_curv / (2 * np.cos(np.radians( default_min_angle ) / 2)) + 0.05 # make sure that highest angle is 60 + add epsilon for corner cases
    for k in range(amount_of_pcl):
        if angle!=0 or radius!=0:
            a=b=c=d=e=k1=k2 = 0
            if angle>0:
                if label==1 or edge==1:
                    # make sure the height of the pyramid is not too low and not too high (using the connection between R and h and the angles of the tip of the pyramid)
                    # pythagoras thm on height of pyr and height of face:  h^2 + 0.25R^2 = 0.75R^2 * 1 / (tan(alpha/2))^2; tan(rise_aang) = h / R
                    # because parabola has curvature 2*a we enforce same values for rise of pyramid at 0.5
                    int_K_const = ((max_curv + 10e-6) ** 2 / (2 * np.pi))
                    boundary_rise_angle_rad = np.arctan( 0.25 * boundary)
                    max_curv_rise_angle_rad = np.arctan( 0.25 * max_curv)
                    min_rise_angle_default = (np.pi - np.radians(default_max_angle)) / 2
                    max_rise_angle_default = (np.pi - np.radians(default_min_angle)) / 2
                    min_rise_angle = max(boundary_rise_angle_rad, min_rise_angle_default)
                    max_rise_angle = min(max_curv_rise_angle_rad, max_rise_angle_default)
                    max_pyr_angle = 2 * (np.arctan(1 / ((0.25 + np.tan(min_rise_angle)**2 ) / 0.75 )))
                    min_pyr_angle = 2 * (np.arctan(1 / ((0.25 + np.tan(max_rise_angle)**2 ) / 0.75 )))

                    if min_pyr_angle < 0:
                        print("fix tis pyramid!!")
                        exit(-1)

                    angle_rad = np.random.uniform(min_pyr_angle, max_pyr_angle)
                    cur_gauss_curv = (2 * np.pi - angle_rad * 3) * int_K_const
                    cur_curve = np.sqrt(cur_gauss_curv)
                    k1 = k2 = cur_curve
                    angle = np.degrees(angle_rad)
                if label==2 or edge==2:
                    max_angle_rad = np.clip(2 * np.arccos(boundary / (2 * constant)), np.radians(default_min_angle), np.radians(default_max_angle))
                    min_angle_rad = np.clip(2 * np.arccos(max_curv / (2 * constant)), np.radians(default_min_angle), np.radians(default_max_angle))
                    angle_rad = np.random.uniform(min_angle_rad, max_angle_rad)
                    cur_curve = constant * (2 * np.cos(angle_rad / 2))
                    k1 = cur_curve
                    k2 = 0
                    angle = np.degrees(angle_rad)
            if radius > 0:
                cur_curve = np.random.uniform(boundary ,max_curv)
                radius = 1 / cur_curve
                if label==1 or edge==1:
                    k1 = k2 = cur_curve
                if label==2 or edge==2:
                    k1 = cur_curve
                    k2 = 0
        else:
            a, b, c, d, e, _, k1, k2 = createFunction(gaussian_curv=gaussian_curv, mean_curv=mean_curv, boundary=boundary, epsilon=epsilon, max_curv=max_curv, size=size)
        dataset_name = f"point_cloud_{counter + k}"
        if update == True:
            del point_clouds_group[dataset_name]
        point_cloud = point_clouds_group.create_dataset(dataset_name, data=np.array([0, 0, 0]).reshape(1, 3))
        point_cloud.attrs['a'] = a
        point_cloud.attrs['b'] = b
        point_cloud.attrs['c'] = c
        point_cloud.attrs['d'] = d
        point_cloud.attrs['e'] = e
        point_cloud.attrs['k1'] = k1
        point_cloud.attrs['k2'] = k2
        point_cloud.attrs['angle'] = angle
        point_cloud.attrs['radius'] = radius
        point_cloud.attrs['class'] = label
        point_cloud.attrs['edge'] = edge

def createFunction(gaussian_curv, mean_curv, max_curv, size, boundary=1, epsilon=0.25):
    if gaussian_curv==1 and mean_curv==0:
        raise ValueError("gaussian_curv==1 and mean_curv==0 is impossible")
    okFunc = False
    count = 0
    while not okFunc:
        okFunc=True
        count += 1
        # a, b, c, d, e = np.random.uniform(-1.3, 1.3, 5)
        a, b, c, d, e = np.random.uniform(-size, size, 5)
        K = (4*(a*b)-((c**2))) / ((1 + d**2 + e**2)**2)
        H = (a*(1 + e**2)-d*e*c +b*(1 + d**2)) / ( ( (d**2) + (e**2) + 1 )**1.5)

        discriminant = H ** 2 - K
        k1 = H + np.sqrt(discriminant)
        k2 = H - np.sqrt(discriminant)


        temp_max = k1 if abs(k1) > abs(k2) else k2
        temp_min = k1 if abs(k1) < abs(k2) else k2

        # Not to steep
        if abs(temp_max) > max_curv:
            okFunc = False
            continue

        # zero gaussian curve --> either plane or ridge/valley
        if gaussian_curv==0:
            if (abs(temp_min) > epsilon):
                okFunc=False
                continue
            # positive mean curv
            if mean_curv == 0:
                if (abs(temp_max) > epsilon):
                    okFunc = False
                    continue
            # positive mean curv
            if mean_curv == 1:
                if temp_max < (boundary):
                    okFunc = False
                    continue
            # negative mean curv
            if mean_curv == -1:
                if temp_max > -(boundary):
                    okFunc = False
                    continue

        # non-zero gaussian curve --> either parabola or saddle
        else:
            if (abs(temp_min) < boundary):
                okFunc = False
                continue
            # positive gaussian curv
            if gaussian_curv==1:
                if K < 0:
                    okFunc=False
                    continue
            # negative gaussian curv
            if gaussian_curv==-1:
                if K > 0:
                    okFunc=False
                    continue

    return a, b, c, d, e, count , k1, k2

if __name__ == '__main__':
    createDataSetOld()
