# Code adapted from Andreas Maisinger (TU Dortmund University)

import numpy as np

try:
    import torch
except ImportError as err:  # pragma: no cover
    raise ModuleNotFoundError(
        "You need torch to use the ppdisks module of radiosim!"
        'Use uv pip install "radiosim[torch]" to install the necessary packages!'
    ) from err

from astropy.convolution import Gaussian2DKernel
from torch.fft import fft2, ifft2
from torchvision.transforms.functional import resize
from tqdm.autonotebook import tqdm

__all__ = ["create_proto", "generate_proto_set"]


def generate_proto_set(
    img_size: int,
    size: int,
    alpha_range: tuple = (0, 180),
    ratio_range: tuple = (3, 15),
    size_ratio_range: tuple = (0.1, 1),
    device: str = "cpu",
    seed: int = 1337,
    verbose: bool = False,
):
    """
    Generates a Set of simulated protoplanetary disk images

    Parameters
    ----------
    img_size: int
        The size of the image (height and width)
    size: int
        The amount of images to simulate
    alpha_range: tuple of float, optional
        The range of values of the inclination of the
        disk in degrees. Default: (0, 180)
    ratio_range: tuple of float, optional
        The range of values of the ratio the minor axis should
        be of the major axis :math:`a / b`. Default: (3, 15)
    size_ratio_range: tuple of float, optional
        The range of values of the ratio of the image
        size the disk should take up. Default: (0.1, 1)
    device: str, optional
        The name of the device to run the simulations
        on (with torch). Default: 'cpu'
    seed: int, optional
        The seed for the random generator. Default: 1337
    verbose: bool, optional
        Whether to show a progress bar or not. Default: False

    Returns
    -------
    protos: list of torch.tensor
        List of the simulated images

    params: dict
        The used simulation parameters

    """

    np.random.seed(seed)

    bundles = tqdm(range(size)) if verbose else range(size)

    outputs = [
        create_proto(
            img_size,
            alpha=np.random.randint(*alpha_range),
            ratio=np.random.randint(*ratio_range) / np.random.randint(*ratio_range),
            size_ratio=np.random.uniform(*size_ratio_range),
            device=device,
        )
        for i in bundles
    ]

    protos = [out.cpu().numpy() for out in outputs]
    params = dict(
        seed=seed,
        N=size,
        img_size=img_size,
        alpha_range=alpha_range,
        ratio_range=ratio_range,
        size_ratio_range=size_ratio_range,
    )

    return protos, params


def create_proto(
    img_size: int,
    alpha: float,
    ratio: float,
    size_ratio: float,
    device: str = "cpu",
    seed: int = None,
):
    """
    Generates a simulated protoplanetary disk image

    Parameters
    ----------
    img_size: int
        The size of the image (height and width)
    alpha: float
        The inclination angle of the disk in degrees
    ratio: float
        The ratio the minor axis should be of the major axis a / b
    size_ratio: float
        The ratio of the image size the disk should take up
    device: str, optional
        The name of the device to run the simulations on (with torch)
    seed: int, optional
        The seed for the random generator. If set to ``None``, the seed is ignored.

    Returns
    -------
    proto: torch.tensor
        The simulated image of the protoplanetary disk

    """

    if seed is not None:
        np.random.seed(seed)

    fix_size = 900
    center = (fix_size // 2, fix_size // 2)
    m, n = _get_m_n()
    core, d_core = _sim_core(fix_size, alpha, center, ratio, device=device)
    sc = np.random.randint(5, 10)
    core = _smooth(sc, sc, core, device=device) * _get_scale(d_core, m, n)
    num_circ = int(np.random.normal(3, 0.5) // 1)

    els_full = torch.zeros((900, 900), device=device)
    ds = []
    a_max = d_core

    for i in range(num_circ):
        eli, di, maxlength = _sim_ring(
            fix_size, alpha, center, ratio, m, n, (i + 1) * 30, device=device
        )
        a_max = np.max([maxlength, a_max])
        els_full += eli
        ds.append(di)

    proto = core + els_full
    proto = torch.abs(proto)

    img_size2 = int(fix_size // 2)

    a_max = np.min([img_size2, a_max])

    proto_snip = proto[
        int(img_size2 - a_max) : int(img_size2 + a_max),
        int(img_size2 - a_max) : int(img_size2 + a_max),
    ]

    if proto_snip.max() < 1e-10:
        return create_proto(img_size, alpha, ratio, size_ratio)

    new_size = int(size_ratio * img_size)
    new_size = new_size if new_size % 2 == 0 else new_size - 1

    proto_snip = proto_snip[None, :, :]

    proto_snip = resize(proto_snip, (new_size, new_size))
    proto_snip = proto_snip[0]

    proto_snip_size2 = int(proto_snip.shape[0] // 2)

    img_size2 = int(img_size // 2)

    proto = torch.ones((img_size, img_size), device=device) * 1e-10
    proto[
        img_size2 - proto_snip_size2 : img_size2 + proto_snip_size2,
        img_size2 - proto_snip_size2 : img_size2 + proto_snip_size2,
    ] += proto_snip

    del proto_snip
    return proto.cpu()


def _create_ellipse(
    img_size: int,
    a: float,
    b: float,
    center: tuple,
    alpha: float,
    device: str = "cpu",
):
    """
    Generates an ellipse

    Parameters
    ----------
    img_size: int
        The size of the image (height and width)
    a: float
        The major axis of the ellipse
    b: float
        The minor axis of the ellipse
    center: tuple of float
        The center of the ellipse
    alpha: float
        The rotation angle of the ellipse
    device: str, optional
        The name of the device to run the simulations on (with torch)

    Returns
    -------
    ellipse: torch.tensor
        The image of the ellipse
    average_distance: float
        The average distance from the center of the ellipse
    """

    x_0 = center[0]
    y_0 = center[1]

    xx, yy = torch.from_numpy(np.mgrid[:img_size, :img_size]).to(device)
    term1 = (
        ((xx - x_0) * np.cos(np.deg2rad(alpha)))
        + ((yy - y_0) * np.sin(np.deg2rad(alpha)))
    ) ** 2
    term2 = (
        ((xx - x_0) * np.sin(np.deg2rad(alpha)))
        - ((yy - y_0) * np.cos(np.deg2rad(alpha)))
    ) ** 2
    ellipse = ((term1 / a**2) + (term2 / b**2)) <= 1
    avarage_distance = (a + b) / 2
    return ellipse, avarage_distance


def _substract(
    ellipse1: torch.tensor, ellipse2: torch.tensor, distance1: float, distance2: float
):
    """
    Substracts two ellipses from eachother

    Parameters
    ----------
    ellipse1: torch.tensor
    The first ellipse

    ellipse2: torch.tensor
    The second ellipse

    distance1: float
    The average distance from the center of the first ellipse

    distance1: float
    The average distance from the center of the second ellipse

    Returns
    -------
    subtracted: torch.tensor
    The result of the subtraction

    average_distance: float
    The average distance from the center of the ellipse resulting from the subtraction

    """

    subtracted = ellipse1 ^ ellipse2
    average_distance = (distance1 + distance2) / 2
    return subtracted, average_distance


def _sim_core(
    img_size: int, alpha: float, center: tuple, ratio: float, device: str = "cpu"
):
    """
    Simulates the core of the protoplanetary disk

    img_size: int
    The size of the image (height and width)

    alpha: float
    The rotation angle of the core

    center: tuple of float
    The center of the core

    ratio: float
    The ratio the minor axis should be of the major axis a / b

    device: str, optional
    The name of the device to run the simulations on (with torch)

    Returns
    -------
    core: torch.tensor
    The image of the core

    d_core: float
    The average distance from the center of the core

    """

    a = np.float64(np.random.randint(3, 15))
    core, d_core = _create_ellipse(img_size, a, a / ratio, center, alpha, device=device)
    return core, d_core


def _sim_ring(
    img_size: int,
    alpha: float,
    center: tuple,
    ratio: float,
    m: float,
    n: float,
    initial_a: float = 0,
    device: str = "cpu",
):
    """
    Simulates a ring of the protoplanetary disk

    img_size: int
    The size of the image (height and width)

    alpha: float
    The rotation angle of the core

    center: tuple of float
    The center of the core

    ratio: float
    The ratio the minor axis should be of the major axis a / b

    m: float
    Empirical value determining the scale of the ring

    n: float
    Empirical value determining the scale of the ring

    initial_a: float, optional
    The initial size of the major axis

    device: str, optional
    The name of the device to run the simulations on (with torch)

    Returns
    -------
    ring: torch.tensor
    The image of the ring

    d_ring: float
    The average distance from the center of the ring

    max_value: float
    Largest extent of the ring

    """

    a = np.random.randint(10, 35) + initial_a
    exp = np.random.randint(5, 10)

    ring_a, d_ring_a = _create_ellipse(
        img_size, a, a / ratio, center, alpha, device=device
    )

    ring_b, d_ring_b = _create_ellipse(
        img_size, a + exp, (a + exp) / ratio, center, alpha, device=device
    )
    ring, d_ring = _substract(ring_b, ring_a, d_ring_b, d_ring_a)

    s = _get_smooth_factor(a)
    ring = _smooth(s, s, ring, device=device) * _get_scale(d_ring, m, n)

    max_value = np.max([a, a + exp, a / ratio, (a + exp) / ratio])

    return ring, d_ring, max_value + 6 * s


def _get_smooth_factor(a: float):
    """
    Empirically determines the standard deviation for the gaussian smoothing.
    """

    if a < 15:
        return 2
    if a < 20:
        return 4
    if a < 30:
        return 6
    if a < 40:
        return 9
    if a >= 40:
        return 12


# from https://stackoverflow.com/a/47979802
def _convolve2d(x: torch.tensor, y: torch.tensor):
    """
    Convolves two 2-dimensional images with eachother

    Parameters
    ----------
    x: torch.tensor
    The first image

    y: torch.tensor
    The second image

    Returns
    -------
    cc: torch.tensor
    The convolved image

    """

    fr = fft2(x)
    fr2 = fft2(torch.flipud(torch.fliplr(y)))
    m, n = fr.shape
    cc = torch.real(ifft2(fr * fr2))
    cc = torch.roll(cc, -int(m / 2) + 1, dims=0)
    cc = torch.roll(cc, -int(n / 2) + 1, dims=1)
    return cc


def _smooth(x: float, y: float, c: torch.tensor, device: str = "cpu"):
    """
    Smooths the input image by convolving it with a gaussian 2d-kernel

    Parameters
    ----------
    x: float
    The standard deviation in x-direction

    y: float
    The standard deviation in y-direction

    c: torch.tensor
    The image to smooth

    device: str, optional
    The name of the device to run the simulations on (with torch)

    """

    gauss_kernel = torch.from_numpy(
        Gaussian2DKernel(
            x_stddev=x, y_stddev=y, x_size=c.shape[0], y_size=c.shape[1]
        ).array
    ).to(device)
    return _convolve2d(c, gauss_kernel)


def _e(x: float, m: float, n: float):
    """
    Computes the value of an exponential function exp(m * x + n)
    """

    return np.exp(m * x + n)


def _get_m_n():
    """
    Get the empirical values for m and n from JSON files
    """

    results = {
        "average_1": {"m": -0.08908080536281453, "n": -4.808388735775005},
        "average_2": {"m": -0.04053517757866287, "n": -5.718461569072661},
        "average_3": {"m": -0.023989028595164514, "n": -6.22105762345096},
        "average_4": {"m": -0.011198837787046627, "n": -6.701733964123176},
    }

    id = np.random.choice(list(results.keys()))
    return results[id]["m"], results[id]["n"]


def _get_scale(distance: float, m: float, n: float, noise_scale: float = 0.05):
    """
    Get the empirical standard deviation for the smoothing
    """

    return _e(distance, m, n) + np.random.normal(0, _e(distance, m, n) * noise_scale)
