import math
import random

import numpy as np

import math
import torch

def circle_intersection(x0, y0, r0, x1, y1, r1):
    # circle 1: (x0, y0), radius r0
    # circle 2: (x1, y1), radius r1

    d=math.sqrt((x1-x0)**2 + (y1-y0)**2)
    
    # non intersecting
    if d > r0 + r1 :
        return None
    # One circle within other
    if d < abs(r0-r1):
        return None
    # coincident circles
    if d == 0 and r0 == r1:
        return None
    else:
        a=(r0**2-r1**2+d**2)/(2*d)
        h=math.sqrt(r0**2-a**2)
        x2=x0+a*(x1-x0)/d   
        y2=y0+a*(y1-y0)/d   
        x3=x2+h*(y1-y0)/d     
        y3=y2-h*(x1-x0)/d 

        x4=x2-h*(y1-y0)/d
        y4=y2+h*(x1-x0)/d
        
        return (np.array([x3, y3]), np.array([x4, y4]))


class MirrorKeypoints:
    def __call__(self, sample):    
        if random.random() > 0.5:
            return sample
        # flip the keypoints tensor
        sample = 1 -  sample

        return sample

class Z_augmentation:

    def __init__(self, hand_side="left"):
        self.hand_side = hand_side

    def new_wrist(self, sample, hand_side="left", new_wrist=None):
        if hand_side == "left":
            wrist = sample[30:32]
            shoulder = sample[22:24]
            elbow = sample[26:28]
        else:
            wrist = sample[32:34]
            shoulder = sample[24:26]
            elbow = sample[28:30]

        # calculate the length of the shoulder to elbow using math package
        shoulder_elbow_length = math.sqrt((shoulder[0] - elbow[0])**2 + (shoulder[1] - elbow[1])**2)
        # calculate the length of the wrist to elbow using math package
        wrist_elbow_length = math.sqrt((wrist[0] - elbow[0])**2 + (wrist[1] - elbow[1])**2)

        if shoulder_elbow_length == 0 or wrist_elbow_length == 0:
            return sample, None

        first_time = True
        new_loc = False
        while not new_loc:

            if new_wrist is None or not first_time:
                # get random new wrist point that is not too far from the elbow
                new_wrist = [random.uniform(elbow[0] - 0.3, elbow[0] + 0.3), random.uniform(elbow[1] - 0.3, elbow[1] + 0.3)]

            # get intersection points of the circles
            c = circle_intersection(shoulder[0], shoulder[1], shoulder_elbow_length, new_wrist[0], new_wrist[1], wrist_elbow_length)
            if c is not None:
                (i1, i2) = c
                new_loc = True
            first_time = False

        # get the point that is below the hand
        if i1[1] > i2[1]:
            new_elbow = i1
        else:
            new_elbow = i2
        # new_elbow to shape (2,1)
        new_elbow = np.array(new_elbow)
        new_wrist = np.array(new_wrist)

        # replace the keypoints in the sample
        if hand_side == "left":
            sample[26:28] = new_elbow
            sample[30:32] = new_wrist
        else:
            sample[28:30] = new_elbow
            sample[32:34] = new_wrist
        return sample, new_wrist
    
    def __call__(self, samples):
        # transform each sample in the batch
        t_new = []

        t = samples.numpy()
        new_wrist = None
        for t_i in t:
            # if new_wrist is None:
            #     new_t, w = self.new_wrist(t_i.reshape(-1), self.hand_side)
            #     new_wrist = w
            # else:
            new_t, _ = self.new_wrist(t_i.reshape(-1), self.hand_side)
            # reshape back to 2 dimensions
            t_new.append(new_t.reshape(-1, 2))
        return torch.tensor(np.array(t_new))
    
# augmentation to add little randow noise to the keypoints
class NoiseAugmentation:
    def __init__(self, noise=0.05):
        self.noise = noise

    def __call__(self, sample):
        # add noise to the keypoints
        sample = sample + torch.randn(sample.shape) * self.noise
        return sample

# augmentation to rotate all keypoints around 0,0
class RotateAugmentation:
    def __call__(self, sample):
        # generate a random angle between -13 and 13 degrees
        angle_max = 13.0
        angle = math.radians(random.uniform(a=-angle_max, b=angle_max))
        # rotate the keypoints around 0.0
        new_sample = sample
        new_sample[:, :, 0] = sample[:, :, 0]*math.cos(angle) - sample[:, :, 1]*math.sin(angle)
        new_sample[:, :, 1] = sample[:, :, 0]*math.sin(angle) + sample[:, :, 1]*math.cos(angle)
        return new_sample