import math
import matplotlib.pyplot as plt
import numpy as np
from tokamak.formats import geqdsk


def generate_ellipse(k, h, r, a, b):
    ellipse = []
    step = math.pi / 1000
    for t in np.arange(0.0, 2 * math.pi, step):
        x = h + a * r * math.cos(t)
        y = k + b * r * math.sin(t)
        ellipse.append((x, y))
    return ellipse


def get_middle(p1, p2):
    return (p1[0] + p2[0]) / 2, (p1[1] + p2[1]) / 2


def find_base_points(curve):
    angles = []
    N = len(curve)
    min_points_dst1 = N // 10
    min_points_dst2 = N // 10

    for i, p in enumerate(curve):
        # three continuous points
        point_1 = curve[i - 1]
        point_2 = curve[i]
        point_3 = curve[(i + 1) % N]
        # 2 vectors to find angle
        vec_1 = (point_1[0] - point_2[0], point_1[1] - point_2[1])
        vec_2 = (point_1[0] - point_3[0], point_1[1] - point_3[1])
        # find an angle

        if math.sqrt(vec_1[0] ** 2 + vec_1[1] ** 2) * math.sqrt(vec_2[0] ** 2 + vec_2[1] ** 2) == 0.0:
            a = math.acos(1.0)
        else:
            a = math.acos((vec_1[0] * vec_2[0] + vec_1[1] * vec_2[1]) / (
                    math.sqrt(vec_1[0] ** 2 + vec_1[1] ** 2) * math.sqrt(vec_2[0] ** 2 + vec_2[1] ** 2)))
        angles.append((a, i))  # save angle and point index
    angles.sort(key=lambda a: a[0])

    ind_min = [angles[0][1]]
    # Find minimum, which is far from first minimum point
    for i in range(1, len(angles)):
        ind = angles[i][1]
        dist = N // 2 - abs(abs(ind - ind_min[0]) - N // 2) + 1
        if dist >= min_points_dst2:
            ind_min.append(ind)
            break

    ind_max = [angles[-1][1]]
    # Find maximum, which is far from first maximum point
    for i in range(len(angles) - 1, 0, -1):  # The more is ind, the more is angle
        ind = angles[i][1]
        dist = N // 2 - abs(abs(ind - ind_max[0]) - N // 2) + 1
        if dist >= min_points_dst1:
            ind_max.append(ind)
            break
    indices = ind_min + ind_max
    return indices


def deep_cut(curve, indices, depth):
    n = len(curve)
    for _ in range(depth):
        indices.sort()
        new_indices = []
        for i in range(len(indices)):
            # Add points between existing ones
            new_ind = ((indices[i - 1] + indices[i] + (0 if indices[i - 1] < indices[i] else n)) // 2) % len(curve)
            new_indices.append(new_ind)
        indices += new_indices
    indices.sort()
    points = [(curve[i][0], curve[i][1]) for i in indices]
    return points


def tokamak_split(x, y, center):
    depth_level = 0

    # Delete similar point to first one
    curve = []
    for i in range(len(x)):
        c = (x[i], y[i])
        curve.append(c)

    # del curve[-1]

    indices = find_base_points(curve)
    points = deep_cut(curve, indices, depth_level)
    # Define lines
    # Configure line containing 2 points for each line
    X = []
    Y = []
    for pnt in points:
        X.append(pnt[0])
        X.append(center[0])
        Y.append(pnt[1])
        Y.append(center[1])
    # Define figure 2
    new_curve_x = []
    new_curve_y = []
    for pnt in points:
        mid_pnt = get_middle(pnt, center)
        new_curve_x.append(mid_pnt[0])
        new_curve_y.append(mid_pnt[1])
    # Duplicate first point to make closed figure
    new_curve_x += [new_curve_x[0]]
    new_curve_y += [new_curve_y[0]]

    return [X, Y], [new_curve_x, new_curve_y]


data = geqdsk.read("g035685.00146")

r, z = data["rbdry"], data["zbdry"]

center_x, center_y = data["bcentr"], data["zmid"]

x_0 = center_x
y_0 = center_y  # center

center = (x_0, y_0)
lines, new_curve = tokamak_split(r, z, center)

# paint 4 lines
fig, ax = plt.subplots()
ax.set_xlim(-0.0, 0.75)
ax.set_ylim(-0.5, 0.5)
ax.grid()
ax.set_aspect(1.0)

ax.set(xlabel='r', ylabel='z',
       title='tokamak')
# x = []
# y = []


plt.plot(r, z, 'r')
plt.plot(center[0], center[1], 'bo')
plt.plot(lines[0], lines[1], 'g')
plt.plot(new_curve[0], new_curve[1], 'b-')

# plt.plot(r, z, 'b-')

plt.savefig("four")
plt.show()