from openmc_plasma_source import TokamakSource
from openmc import Source
import numpy as np

import pytest
from hypothesis import given, settings, assume, strategies as st


@st.composite
def tokamak_source_strategy(draw, return_dict=False):
    """Defines a hypothesis strategy that automatically generates a TokamakSource.
    When passed `return_dict=True`, only a dict of the args is passed.
    Geometry attributes are varied, while plasma attributes are fixed.
    """
    # Used to avoid generation of inappropriate float values
    finites = {"allow_nan": False, "allow_infinity": False, "allow_subnormal": False}

    # Specify the base strategies for each geometry input
    minor_radius = draw(st.floats(min_value=0.0, max_value=100.0, **finites))
    major_radius = draw(st.floats(min_value=0.0, max_value=100.0, **finites))
    pedestal_radius = draw(st.floats(min_value=0.0, max_value=100.0, **finites))
    elongation = draw(st.floats(min_value=1.0, max_value=10.0, **finites))
    triangularity = draw(st.floats(min_value=-1.0, max_value=1.0, **finites))
    shafranov_factor = draw(st.floats(min_value=0.0, max_value=1.0, **finites))

    # Specify requirements that must be satisfied for a valid tokamak
    assume(major_radius > minor_radius)
    assume(minor_radius > pedestal_radius)
    assume(minor_radius > shafranov_factor)

    args_dict = {
        "elongation": elongation,
        "triangularity": triangularity,
        "major_radius": major_radius,
        "minor_radius": minor_radius,
        "pedestal_radius": pedestal_radius,
        "shafranov_factor": shafranov_factor,
        "ion_density_centre": 1.09e20,
        "ion_density_peaking_factor": 1,
        "ion_density_pedestal": 1.09e20,
        "ion_density_separatrix": 3e19,
        "ion_temperature_centre": 45.9,
        "ion_temperature_peaking_factor": 8.06,
        "ion_temperature_pedestal": 6.09,
        "ion_temperature_separatrix": 0.1,
        "mode": "H",
        "ion_temperature_beta": 6,
    }

    return args_dict if return_dict else TokamakSource(**args_dict)


@given(tokamak_source=tokamak_source_strategy())
@settings(max_examples=1, deadline=None)
def test_creation(tokamak_source):
    """Tests that the sources generated by TokamakSource are of type openmc.Source"""
    for source in tokamak_source.sources:
        assert isinstance(source, Source)


@given(tokamak_source=tokamak_source_strategy())
@settings(max_examples=100, deadline=None)
def test_strengths_are_normalised(tokamak_source):
    """Tests that the sum of the strengths attribute is equal to"""
    assert pytest.approx(sum(tokamak_source.strengths), 1)


@given(tokamak_source=tokamak_source_strategy())
@settings(max_examples=100, deadline=None)
def test_source_locations_are_within_correct_range(tokamak_source):
    """Tests that each source has RZ locations within the expected range.
    As the function converting (a,alpha) coordinates to (R,Z) is not bijective, we
    cannot convert back to validate each individual point. However, we can determine
    whether the generated points are contained within the shell of the last closed
    magnetic surface.
    See "Tokamak D-T neutron source models for different plasma physics confinement
    modes", C. Fausser et al., Fusion Engineering and Design, 2012 for more info.
    """
    R_0 = tokamak_source.major_radius
    A = tokamak_source.minor_radius
    El = tokamak_source.elongation
    delta = tokamak_source.triangularity

    def get_R_on_LCMS(alpha):
        """Gets R on the last closed magnetic surface for a given alpha"""
        return R_0 + A * np.cos(alpha + delta * np.sin(alpha))

    approx_lt = lambda x, y: x < y or np.isclose(x, y)
    approx_gt = lambda x, y: x > y or np.isclose(x, y)

    for source in tokamak_source.sources:
        R, Z = source.space.r.x[0], source.space.z.x[0]
        # First test that the point is contained with a simple box with
        # lower left (r_min,-z_max) and upper right (r_max,z_max)
        assert approx_gt(R, R_0 - A)
        assert approx_lt(R, R_0 + A)
        assert approx_lt(abs(Z), A * El)
        # For a given Z, we can determine the two values of alpha where
        # where a = minor_radius, and from there determine the upper and
        # lower bounds for R.
        alpha_1 = np.arcsin(abs(Z) / (El * A))
        alpha_2 = np.pi - alpha_1
        R_max, R_min = get_R_on_LCMS(alpha_1), get_R_on_LCMS(alpha_2)
        assert approx_lt(R_max, R_0 + A)
        assert approx_gt(R_min, R_0 - A)
        assert approx_lt(R, R_max)
        assert approx_gt(R, R_min)


@given(tokamak_args_dict=tokamak_source_strategy(return_dict=True))
@settings(max_examples=1, deadline=None)
def test_angles(tokamak_args_dict):
    """Checks that custom angles can be set"""
    tokamak_args_dict["angles"] = (0, 1)
    tokamak_source = TokamakSource(**tokamak_args_dict)
    assert tokamak_source.angles == (0, 1)
    for source in tokamak_source.sources:
        assert (source.space.phi.a, source.space.phi.b) == (0, 1)
