#!/usr/bin/env python3
import matplotlib.pyplot as plt
import numpy as np
from dataclasses import dataclass


@dataclass
class Defect:
    x: float
    y: float
    s: float
    beta: float


def lc(x, y, defects):
    dx = 0
    dy = 0
    for d in defects:
        dist = np.sqrt((x - d.x)**2 + (y - d.y)**2)
        angle = np.atan2(y - d.y, x - d.x)
        xdir = np.cos(angle * d.s + d.beta) / dist
        ydir = np.sin(angle * d.s + d.beta) / dist

        dx += xdir
        dy += ydir
    return dx, dy

def draw_field(n, D, filename=None, title=None, show=True):
    xrange = np.arange(-n,n,0.1)
    yrange = np.arange(-n,n,0.1)
    X,Y = np.meshgrid(xrange,yrange)

    Ex, Ey = lc(X, Y, D)

    plt.figure(figsize=(10, 10))
    for d in D:
        plt.plot(d.x, d.y, '-og')
    plt.streamplot(X, Y, Ex, Ey, density=1.4, linewidth=None, color='#A23BEC')

    plt.grid()
    if title != None:
        plt.title(title)
    if filename != None:
        plt.savefig(filename, bbox_inches='tight')
    if show == True:
        plt.show()

n = 20
defects = [
    ( 1, 0),
    (-1, 0),
    ( 1, np.pi/2),
    ( 0.5, 0),
    (-0.5, 0),
    (-1.5, 0),
    ( 1.5, 0),
    ( 2, 0),
    (-2, 0),
    ( 2.5, 0),
    (-2.5, 0),
    ( 3, 0),
    (-3, 0),
    ( 3.5, 0),
    (-3.5, 0),
    ( 4, 0),
    (-4, 0),
]
for s, beta in defects:
    D = []
    D.append(Defect(0, 0, s, beta))
    print(D)
    draw_field(n, D,
            title=f'Director field, s = {s}, β = {beta:.2f}',
            filename=f'director_field,s={s:+.1f},beta={beta:.2f}.png',
            show=False
            )
 
# D = []
# D.append(Defect(0, 0, -1.5, 0))
# print(D)
# draw_field(n, D)