#!/usr/bin/env python
# -*- coding: utf-8 -*-

__all__ = [
    "parallel_beam_c",
    "parallel_beam_p",
    "point_source_c",
    "point_source_p",
    "point_source_r",
]
"""Module with functions to define ray sources"""


from numpy import sin, cos, pi
from numpy.random import normal

from .ray import Ray


def parallel_beam_c(
    origin=(0.0, 0.0, 0.0),
    direction=(0.0, 0.0, 0.0),
    size=(1.0, 1.0),
    num_rays=(10, 10),
    wavelength=0.58929,
    label="",
    draw_color=None,
):
    """Cartesian grid parallel beam

    This function creates a parallel beam, where the rays are organized in a
    cartesian grid.

    Parameters:


    *origin*
        Tuple with the coordinates of the central ray origin
    *direction*
        Tuple with the rotation of the beam around the XYZ axes.
    *size*
        Tuple with the beam's width and the height.
    *num_rays*
        Tuple (nx,ny) containing the number of rays used to create the beam.
    *label*
        String used to identify the ray source
    *draw_color*
        Color used to represent the rays in plots. Default (None) automatically
        chooses a color based on the wavelength. Otherwise, can be any valid
        matplotlib color descriptor. See :
            https://matplotlib.org/3.3.1/api/colors_api.html
    """

    ret_val = []

    nx, ny = num_rays
    dx, dy = size

    # note modify this to use traits
    dx = float(dx)
    dy = float(dy)

    for ix in range(nx):
        for iy in range(ny):
            x = -dx / 2.0 + dx * ix / (nx - 1)
            y = -dy / 2.0 + dy * iy / (ny - 1)
            ret_val.append(
                Ray(
                    pos=(x, y, 0),
                    dir=(0, 0, 1),
                    wavelength=wavelength,
                    label=label,
                    draw_color=draw_color,
                ).ch_coord_sys_inv(origin, direction)
            )
    return ret_val


def parallel_beam_p(
    origin=(0.0, 0.0, 0.0),
    direction=(0.0, 0.0, 0),
    radius=0.5,
    num_rays=(5, 10),
    wavelength=0.58929,
    label="",
    draw_color=None,
):
    """Polar grid parallel beam

    This function creates a parallel beam, where the rays are organized in a
    polar grid.

    Parameters:


    *origin*
        Tuple with the coordinates of the central ray origin
    *direction*
        Tuple with the rotation of the beam around the XYZ axes.
    *r*
        Beam radious
    *num_rays*
        Tuple (nr,na) containing the number of rays used to create the beam.
    *label*
        String used to identify the ray source
    *draw_color*
        Color used to represent the rays in plots. Default (None) automatically
        chooses a color based on the wavelength. Otherwise, can be any valid
        matplotlib color descriptor. See :
            https://matplotlib.org/3.3.1/api/colors_api.html
    """

    ret_val = [
        Ray(
            pos=(0, 0, 0), dir=(0, 0, 1), wavelength=wavelength, label=label
        ).ch_coord_sys_inv(origin, direction)
    ]
    nr, nt = num_rays
    for r_ in range(1, nr):
        r = radius * float(r_) / (nr - 1)
        for t in range(nt):
            x_ = r * sin(2 * pi * t / nt)
            y_ = r * cos(2 * pi * t / nt)
            ret_val.append(
                Ray(
                    pos=(x_, y_, 0),
                    dir=(0, 0, 1),
                    wavelength=wavelength,
                    label=label,
                    draw_color=draw_color,
                ).ch_coord_sys_inv(origin, direction)
            )
    return ret_val


def point_source_c(
    origin=(0.0, 0.0, 0.0),
    direction=(0.0, 0.0, 0),
    span=(pi / 8, pi / 8),
    num_rays=(10, 10),
    wavelength=0.58929,
    label="",
    draw_color=None,
):
    """Point source, with a cartesian beam distribution

    This function creates a point source, where the rays are organized in a
    cartesian grid.

    Parameters:


    *origin*
        Tuple with the coordinates of the central ray origin
    *direction*
        Tuple with the rotation of the beam around the XYZ axes.
    *span*
        Tuple angular size of the ray pencil.
    *num_rays*
        Tuple (nx,ny) containing the number of rays used to create the beam.
    *label*
        String used to identify the ray source
    *draw_color*
        Color used to represent the rays in plots. Default (None) automatically
        chooses a color based on the wavelength. Otherwise, can be any valid
        matplotlib color descriptor. See :
            https://matplotlib.org/3.3.1/api/colors_api.html
    """
    ret_val = []

    nx, ny = num_rays
    dx, dy = span

    for ix in range(nx):
        for iy in range(ny):
            if nx != 1:
                tx = -dx / 2.0 + dx * ix / (nx - 1)
            else:
                tx = 0.0

            if ny != 1:
                ty = -dy / 2.0 + dy * iy / (ny - 1)
            else:
                ty = 0.0
            temp_ray = Ray(
                pos=(0, 0, 0),
                dir=(0, 0, 1),
                wavelength=wavelength,
                label=label,
                draw_color=draw_color,
            ).ch_coord_sys_inv((0, 0, 0), (tx, ty, 0))
            ret_val.append(temp_ray.ch_coord_sys_inv(origin, direction))
    return ret_val


def point_source_p(
    origin=(0.0, 0.0, 0.0),
    direction=(0.0, 0.0, 0),
    span=pi / 8,
    num_rays=(10, 10),
    wavelength=0.58929,
    label="",
    draw_color=None,
):
    """Point source, with a polar beam distribution

    This function creates a point source, where the rays are organized in a
    polar grid.

    Parameters:


    *origin*
        Tuple with the coordinates of the central ray origin
    *direction*
        Tuple with the rotation of the beam around the XYZ axes.
    *span*
        Tuple angular size of the ray pencil.
    *num_rays*
        Tuple (nr,na) containing the number of rays used to create the beam
        # TODO Please be more descriptive. It seems really hard to know how much rays we'll get in advance. What is nr?
        # TODO What is na?
    *label*
        String used to identify the ray source
    *draw_color*
        Color used to represent the rays in plots. Default (None) automatically
        chooses a color based on the wavelength. Otherwise, can be any valid
        matplotlib color descriptor. See :
            https://matplotlib.org/3.3.1/api/colors_api.html
    """

    ret_val = [
        Ray(
            pos=(0, 0, 0), dir=(0, 0, 1), wavelength=wavelength, label=label
        ).ch_coord_sys_inv(origin, direction)
    ]

    nr, nt = num_rays

    for r_ in range(1, nr):
        r = span * float(r_) / nr
        temp_ray = Ray(
            pos=(0, 0, 0),
            dir=(0, 0, 1),
            wavelength=wavelength,
            label=label,
            draw_color=draw_color,
        ).ch_coord_sys_inv((0, 0, 0), (r, 0, 0))
        for t in range(nt):
            tz = 2 * pi * t / nt
            temp_ray1 = temp_ray.ch_coord_sys_inv((0, 0, 0), (0, 0, tz))
            ret_val.append(temp_ray1.ch_coord_sys_inv(origin, direction))
    return ret_val


def point_source_r(
    origin=(0.0, 0.0, 0.0),
    direction=(0.0, 0.0, 0),
    span=pi / 8,
    num_rays=100,
    wavelength=0.58929,
    label="",
    draw_color=None,
):
    """Point source, with a ranrom beam distribution

    This function creates a point source, where the rays are organized in a
    random grid.

    Parameters:


    *origin*
        Tuple with the coordinates of the central ray origin
    *direction*
        Tuple with the rotation of the beam around the XYZ axes.
    *span*
        Tuple angular size of the ray pencil.
    *num_rays*
        Number of rays used to create the beam
    *label*
        String used to identify the ray source
    *draw_color*
        Color used to represent the rays in plots. Default (None) automatically
        chooses a color based on the wavelength. Otherwise, can be any valid
        matplotlib color descriptor. See :
            https://matplotlib.org/3.3.1/api/colors_api.html
    """

    ret_val = []

    for n_ in range(num_rays):
        rx = normal(0, span)
        ry = normal(0, span)
        temp_ray = Ray(
            pos=(0, 0, 0),
            dir=(0, 0, 1),
            wavelength=wavelength,
            label=label,
            draw_color=draw_color,
        ).ch_coord_sys_inv((0, 0, 0), (rx, ry, 0))
        ret_val.append(temp_ray.ch_coord_sys_inv(origin, direction))
    return ret_val
