'''
Created on 12-Jul-2021

@author: rishi_priyansh
'''
import random
import inspect
import traceback
import re

import numpy as np
import matplotlib.pyplot as plt
from sympy import Symbol

from xv.km.util import latex_fraction


class FieldOfForce:
    def __init__(self, pos, charge, **kwargs):
        self.pos = np.asarray(pos)
        self.charge = charge
        
    def compute_field(self, x, y):
        X, Y = np.meshgrid(x, y)
        u_i = np.hstack((X.ravel()[:, np.newaxis],  Y.ravel()[:, np.newaxis])) - self.pos
        r = np.sqrt((X - self.pos[0])**2 + (Y - self.pos[1])**2)
        field = ((self.charge / r**2).ravel()[:, np.newaxis] * u_i).reshape(X.shape + (2,))
        return field
    
    def compute_potential(self, x, y):
        X, Y = np.meshgrid(x, y)
        r = np.sqrt((X - self.pos[0])**2 + (Y - self.pos[1])**2)
        potential = self.charge / r
        return potential


    def plot_electric_field(self, lim_value = 5):
        x0, y0 = self.pos  
        x = np.linspace(-lim_value + x0, lim_value + x0, 100)
        y = np.linspace(-lim_value + y0, lim_value + y0, 80)
        
        Y, X = np.meshgrid(x, y)

        field1 = self.compute_field(x, y)
        
        fig, ax1 = plt.subplots(nrows=1, figsize=(6, 6))
        ax1.streamplot(x, y, u=field1[:, :, 0], v=field1[:, :, 1])
        ax1.set_title(f"particle with {self.charge} charge");
        ax1.axis('equal')
        
        pass


    def plot_electric_potential(self, lim_value = 5):
        x0, y0 = self.pos  
        x = np.linspace(-lim_value + x0, lim_value + x0, 100)
        y = np.linspace(-lim_value + y0, lim_value + y0, 80)

        pot1 = self.compute_potential(x, y)
        
        fig, ax1 = plt.subplots(nrows=1, figsize=(6, 6))
        map1 = ax1.pcolormesh(x, y, pot1, 
                              vmin=-lim_value * 2, 
                              vmax=lim_value * 2, 
                              shading='auto')
        ax1.set_title(f"particle with {self.charge} charge");
        ax1.axis('equal')
        plt.colorbar(map1, ax=ax1)
        
    def compute_resulting_field(self, particles, lim_value = 5, 
                                x = None, y = None):
        if x is None or y is None:
            x0, y0 = self.pos  
            if x is None:
                x = np.linspace(-lim_value + x0, lim_value + x0, 100)
            if y is None:
                y = np.linspace(-lim_value + y0, lim_value + y0, 80)

        fields = [p.compute_field(x, y) for p in particles]
        total_field = np.zeros_like(fields[0])
        for field in fields:
            total_field += field
        return total_field
        
    def plot_resulting_field(self, other_particles, lim_value = 5, 
                             x = None, y = None):
        if x is None or y is None:
            x0, y0 = self.pos  
            if x is None:
                x = np.linspace(-lim_value + x0, lim_value + x0, 100)
            if y is None:
                y = np.linspace(-lim_value + y0, lim_value + y0, 80)

        particles = [self, *other_particles]
        total_field = self.compute_resulting_field(particles, 
                                                   lim_value = 5,
                                                   x = x, y = y)
        
        plt.streamplot(x, y, total_field[:, :, 0], total_field[:, :, 1])
        plt.xlim(x.min(), x.max())
        plt.ylim(y.min(), y.max());
        pass
        
        
    pass
#class FieldOfForce:
