# -*- coding: utf-8 -*-
import itertools

import torch

from . import utils


class FieldObject(object):
    def potential(x, y, charge, coupling):
        pass


class Ring(FieldObject):
    def __init__(self, radius: float, charge_density: float = 1.0,
                 x0: float = 0.0, y0: float = 0.0):
        super().__init__()
        self.radius = radius
        self.charge_density = charge_density
        self.x0 = x0
        self.y0 = y0
        
    def potential(self, x: torch.Tensor, y: torch.Tensor, charge: float, coupling: float = 1.0):
        """
        Potential generated by the object on (x, y)

        Parameters
        ----------
        x : torch.Tensor
            Position x-coordinate.
        y : torch.Tensor
            Position y-coordinate.
        charge : float
            Charge of particles.
        coupling : float, optional
            Coupling constant. The default is 1.0.

        Returns
        -------
        value : torch.Tensor
            Value of potential.

        """
        r, _ = utils.to_polar(x - self.x0, y - self.y0)
        value = coupling*self.charge_density*charge*utils.circle_phi(r/self.radius)
        return value


class HorizontalLine(FieldObject):
    def __init__(self, y0: float, charge_density: float = 1.0):
        self.charge_density = charge_density
        self.y0 = y0
        
    def potential(self, x: torch.Tensor, y: torch.Tensor, charge: float, coupling: float = 1.0):
        """
        Potential generated by the object on (x, y)

        Parameters
        ----------
        x : torch.Tensor
            Position x-coordinate.
        y : torch.Tensor
            Position y-coordinate.
        charge : float
            Charge of particles.
        coupling : float, optional
            Coupling constant. The default is 1.0.

        Returns
        -------
        value : torch.Tensor
            Value of potential.

        """
        value = -coupling*self.charge_density*charge*utils.torch.log(torch.abs(y - self.y0))
        return value


class VerticalLine(FieldObject):
    def __init__(self, x0, charge_density=1.0):
        self.charge_density = charge_density
        self.x0 = x0
        
    def potential(self, x: torch.Tensor, y: torch.Tensor, charge: float, coupling: float = 1.0):
        """
        Potential generated by the object on (x, y)

        Parameters
        ----------
        x : torch.Tensor
            Position x-coordinate.
        y : torch.Tensor
            Position y-coordinate.
        charge : float
            Charge of particles.
        coupling : float, optional
            Coupling constant. The default is 1.0.

        Returns
        -------
        value : torch.Tensor
            Value of potential.

        """
        value = -coupling*self.charge_density*charge*torch.log(torch.abs(x - self.x0))
        return value
    

class Hash(FieldObject):
    def __init__(self, l, charge_density=1.0, x0=0.0, y0=0.0):
        self.l = l
        self.charge_density = charge_density
        self.x0 = x0
        self.y0 = y0
        self._set_lines()
    
    def potential(self, x: torch.Tensor, y: torch.Tensor, charge: float, coupling: float = 1.0):
        """
        Potential generated by the object on (x, y)

        Parameters
        ----------
        x : torch.Tensor
            Position x-coordinate.
        y : torch.Tensor
            Position y-coordinate.
        charge : float
            Charge of particles.
        coupling : float, optional
            Coupling constant. The default is 1.0.

        Returns
        -------
        value : torch.Tensor
            Value of potential.

        """
        value = self._upper.potential(x, y, charge, coupling) + \
                self._lower.potential(x, y, charge, coupling) + \
                self._left.potential(x, y, charge, coupling) + \
                self._right.potential(x, y, charge, coupling)
        return value
    
    def _set_lines(self):
        self._upper = HorizontalLine(self.y0 + self.l/2, self.charge_density)
        self._lower = HorizontalLine(self.y0 - self.l/2, self.charge_density)
        self._left = VerticalLine(self.x0 + self.l/2, self.charge_density)
        self._right = VerticalLine(self.x0 - self.l/2, self.charge_density)
        
        
class HorizontalFiniteLine(FieldObject):
    def __init__(self, y0, l, x0=0.0, charge_density=1.0):
        self.charge_density = charge_density
        self.y0 = y0
        self.x0 = x0
        self.l = l
        
    def potential(self, x: torch.Tensor, y: torch.Tensor, charge: float, coupling: float = 1.0):
        """
        Potential generated by the object on (x, y)

        Parameters
        ----------
        x : torch.Tensor
            Position x-coordinate.
        y : torch.Tensor
            Position y-coordinate.
        charge : float
            Charge of particles.
        coupling : float, optional
            Coupling constant. The default is 1.0.

        Returns
        -------
        value : torch.Tensor
            Value of potential.

        """
        dx, dy = x - self.x0, y - self.y0
        integral = torch.arcsinh((2*dx + self.l)/(2*torch.abs(dy))) - \
                   torch.arcsinh((2*dx - self.l)/(2*torch.abs(dy)))
        value = coupling*self.charge_density*charge*integral
        return value


class VerticalFiniteLine(FieldObject):
    def __init__(self, x0, l, y0=0, charge_density=1.0):
        self.charge_density = charge_density
        self.x0 = x0
        self.l = l
        self.y0 = y0
        
    def potential(self, x: torch.Tensor, y: torch.Tensor, charge: float, coupling: float = 1.0):
        """
        Potential generated by the object on (x, y)

        Parameters
        ----------
        x : torch.Tensor
            Position x-coordinate.
        y : torch.Tensor
            Position y-coordinate.
        charge : float
            Charge of particles.
        coupling : float, optional
            Coupling constant. The default is 1.0.

        Returns
        -------
        value : torch.Tensor
            Value of potential.

        """
        dx, dy = x - self.x0, y - self.y0
        integral = torch.arcsinh((2*dy + self.l)/(2*torch.abs(dx))) - \
                   torch.arcsinh((2*dy - self.l)/(2*torch.abs(dx)))
        value = coupling*self.charge_density*charge*integral
        return value


class Square(FieldObject):
    def __init__(self, l, charge_density=1.0, x0=0.0, y0=0.0):
        self.l = l
        self.charge_density = charge_density
        self.x0 = x0
        self.y0 = y0
        self._set_lines()
    
    def potential(self, x: torch.Tensor, y: torch.Tensor, charge: float, coupling: float = 1.0):
        """
        Potential generated by the object on (x, y)

        Parameters
        ----------
        x : torch.Tensor
            Position x-coordinate.
        y : torch.Tensor
            Position y-coordinate.
        charge : float
            Charge of particles.
        coupling : float, optional
            Coupling constant. The default is 1.0.

        Returns
        -------
        value : torch.Tensor
            Value of potential.

        """
        value = self._upper.potential(x, y, charge, coupling) + \
                self._lower.potential(x, y, charge, coupling) + \
                self._left.potential(x, y, charge, coupling) + \
                self._right.potential(x, y, charge, coupling)
        return value
    
    def _set_lines(self):
        self._upper = HorizontalFiniteLine(self.y0 + self.l/2, self.l, self.x0,
                                           self.charge_density)
        self._lower = HorizontalFiniteLine(self.y0 - self.l/2, self.l, self.x0,
                                           self.charge_density)
        self._left = VerticalFiniteLine(self.x0 + self.l/2, self.l, self.y0,
                                        self.charge_density)
        self._right = VerticalFiniteLine(self.x0 - self.l/2, self.l, self.y0,
                                         self.charge_density)


class FixedPoints(FieldObject):
    def __init__(self, x0, y0, charge=1.0):
        super().__init__()
        self.x0 = x0 #(m, )
        self.y0 = y0 #(m, )
        self.charge = charge
        #Assert blablabla

    def potential(self, x: torch.Tensor, y: torch.Tensor, charge: float, coupling: float = 1.0):
        """
        Potential generated by the object on (x, y)

        Parameters
        ----------
        x : torch.Tensor
            Position x-coordinate.
        y : torch.Tensor
            Position y-coordinate.
        charge : float
            Charge of particles.
        coupling : float, optional
            Coupling constant. The default is 1.0.

        Returns
        -------
        value : torch.Tensor
            Value of potential.
        """
        #x : (m,)
        #y : (m,)
        #x0 : (n,)
        #y0 : (n,)
        x = x[..., None]
        y = y[..., None]
        x0 = self.x0 #(1, n)
        y0 = self.y0 #(1, n)
        d = torch.sqrt((x - x0)**2 + (y - y0)**2) #(m, n)
        values = coupling*self.charge*charge*torch.sum(1/d, axis=-1) #(m,)
        return values
        #values = coupling*self.charge*charge/d
        #return values
    
    
class PeriodicFixedPoints(FieldObject):
    def __init__(self, x0, y0, charge=1.0, lx=None, ly=None, cx=0.0, cy=0.0):
        super().__init__()
        self.x0 = x0 #(n, )
        self.y0 = y0 #(n, )
        self.lx = lx
        self.ly = ly
        self.cx = cx
        self.cy = cy
        self.charge = charge #(n, )
        self.nper = 1

    def potential(self, x: torch.Tensor, y: torch.Tensor, charge: float, coupling: float = 1.0):
        """
        Potential generated by the object on (x, y)

        Parameters
        ----------
        x : torch.Tensor
            Position x-coordinate.
        y : torch.Tensor
            Position y-coordinate.
        charge : float
            Charge of particles.
        coupling : float, optional
            Coupling constant. The default is 1.0.

        Returns
        -------
        value : torch.Tensor
            Value of potential.
        """
        x = x[..., None]
        y = y[..., None]
        x0 = self.x0 #(1, n)
        y0 = self.y0 #(1, n)
        d = torch.sqrt((x - x0)**2 + (y - y0)**2) #(m, n)
        values = coupling*self.charge*charge*torch.sum(1/d, axis=-1) #(m,)
        xiterator = list(range(-self.nper, self.nper + 1)) if self.lx is not None else [0]
        yiterator = list(range(-self.nper, self.nper + 1)) if self.ly is not None else [0]
        iterator = itertools.product(xiterator, yiterator)
        single_values = [self.single_potential(x, y, n, m, charge, coupling) 
                         for n, m in iterator]
        values = sum(single_values)
        return values

    def single_potential(self, x, y, n, m, charge, coupling):
        lx = self.lx if self.lx is not None else 0.0
        ly = self.ly if self.ly is not None else 0.0
        d = torch.sqrt((x - self.x0 + n*lx)**2 + (y - self.y0 + m*ly)**2) #(m, n)
        values = coupling*self.charge*charge*torch.sum(1/d, axis=-1) #(m,)
        return values
        