import numpy as np
import scipy.io
from pytest import approx, raises
from pycufsm import cutwp, preprocess, helpers, fsm
import pycufsm.examples.example_1 as example_1
from .utils import pspec_context
from .fixtures.e2e_fixtures import *


def mat_file_test(mat_filename):
    mat = scipy.io.loadmat("tests/mat_files/" + mat_filename)
    cufsm_input = helpers.load_mat(mat)

    # Disable cFSM for now in basic loading of mat files
    cufsm_input["GBTcon"]["glob"] = [0]
    cufsm_input["GBTcon"]["dist"] = [0]
    cufsm_input["GBTcon"]["local"] = [0]
    cufsm_input["GBTcon"]["other"] = [0]

    coords = cufsm_input["nodes"][:, 1:3]
    ends = cufsm_input["elements"][:, 1:4]
    sect_props_cutwp = cutwp.prop2(coords, ends)
    m_all_ones = np.ones((len(cufsm_input["lengths"]), 1))
    signature, curve, shapes = fsm.strip(
        props=cufsm_input["props"],
        nodes=cufsm_input["nodes"],
        elements=cufsm_input["elements"],
        lengths=cufsm_input["lengths"],
        springs=cufsm_input["springs"],
        constraints=cufsm_input["constraints"],
        gbt_con=cufsm_input["GBTcon"],
        b_c='S-S',
        m_all=m_all_ones,
        n_eigs=1,
        sect_props=sect_props_cutwp
    )
    lengths = list(cufsm_input["lengths"])
    return lengths, mat["curve"], sect_props_cutwp, signature, curve, shapes


def describe_end_to_end_tests():

    @pspec_context("End-to-End Tests")
    def describe():
        pass

    def context_example_1():

        @pspec_context("Example 1")
        def describe():
            pass

        results = example_1.__main__()

        def it_results_in_correct_solution():
            assert results["X_values"] == [
                0.5, 0.75, 1, 1.25, 1.5, 1.75, 2, 2.25, 2.5, 2.75, 3, 3.25, 3.5, 3.75, 4, 4.25, 4.5,
                4.75, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16, 17, 18, 19, 20, 22, 24, 26, 28, 30,
                32, 34, 36, 38, 40, 42, 44, 46, 48, 50, 52, 54, 56, 58, 60, 66, 72, 78, 84, 90, 96,
                102, 108, 114, 120, 132, 144, 156, 168, 180, 204, 228, 252, 276, 300
            ]

    def context_dsm_3_2_1_c_with_lips_Mx():

        @pspec_context("DSM Guide Jan 2006, Ex 3.2.1: C-Section with Lips (Mx)")
        def describe():
            pass

        lengths, expected_curve, sect_props, signature, curve, _ = mat_file_test("cwlip_Mx.mat")

        expected = {"xcrl": 5, "Mcrl": 0.67, "xcrd": 26.6, "Mcrd": 0.85, "A": 0.880, "Ixx": 10.285}

        def ir_results_in_correct_sect_props():
            assert expected["A"] == approx(sect_props["A"], abs=0.001)
            assert expected["Ixx"] == approx(sect_props["Ixx"], abs=0.001)

        def it_results_in_correct_Mcr():
            assert signature[lengths.index(expected["xcrl"])] == approx(expected["Mcrl"], abs=0.01)
            assert signature[lengths.index(expected["xcrd"])] == approx(expected["Mcrd"], abs=0.01)

        def it_results_in_correct_signature_curve():
            assert np.allclose(expected_curve[:, 1, 0], curve[:, 0], rtol=1.e-4)

    def context_dsm_3_2_1_c_with_lips_P():

        @pspec_context("DSM Guide Jan 2006, Ex 3.2.1: C-Section with Lips (P)")
        def describe():
            pass

        lengths, expected_curve, sect_props, signature, curve, _ = mat_file_test("cwlip_P.mat")

        expected = {
            "xcrl": 6.6,
            "Pcrl": 0.12,
            "xcrd": 28.5,
            "Pcrd": 0.27,
            "A": 0.880,
            "Ixx": 10.285
        }

        def ir_results_in_correct_sect_props():
            assert expected["A"] == approx(sect_props["A"], abs=0.001)
            assert expected["Ixx"] == approx(sect_props["Ixx"], abs=0.001)

        def it_results_in_correct_Pcr():
            assert signature[lengths.index(expected["xcrl"])] == approx(expected["Pcrl"], abs=0.01)
            assert signature[lengths.index(expected["xcrd"])] == approx(expected["Pcrd"], abs=0.01)

        def it_results_in_correct_signature_curve():
            assert np.allclose(expected_curve[:, 1, 0], curve[:, 0], rtol=1.e-4)

    def context_dsm_3_2_2_c_with_lips_modified_Mx():

        @pspec_context("DSM Guide Jan 2006, Ex 3.2.2: Modified C-Section with Lips (Mx)")
        def describe():
            pass

        lengths, expected_curve, sect_props, signature, curve, _ = mat_file_test(
            "cwlip_modified_Mx.mat"
        )

        expected = {
            "xcrl": 2.7,
            "Mcrl": 1.40,
            "xcrd": 30.5,
            "Mcrd": 0.98,
            "A": 0.933,
            "Ixx": 10.818
        }

        def ir_results_in_correct_sect_props():
            assert expected["A"] == approx(sect_props["A"], abs=0.001)
            assert expected["Ixx"] == approx(sect_props["Ixx"], abs=0.001)

        def it_results_in_correct_Mcr():
            assert signature[lengths.index(expected["xcrl"])] == approx(expected["Mcrl"], abs=0.01)
            assert signature[lengths.index(expected["xcrd"])] == approx(expected["Mcrd"], abs=0.01)

        def it_results_in_correct_signature_curve():
            assert np.allclose(expected_curve[:, 1, 0], curve[:, 0], rtol=1.e-3)

    def context_dsm_3_2_2_c_with_lips_modified_P():

        @pspec_context("DSM Guide Jan 2006, Ex 3.2.2: Modified C-Section with Lips (P)")
        def describe():
            pass

        lengths, expected_curve, sect_props, signature, curve, _ = mat_file_test(
            "cwlip_modified_P.mat"
        )

        expected = {
            "xcrl": 11.5,
            "Pcrl": 0.27,
            "xcrd": 32.7,
            "Pcrd": 0.32,
            "A": 0.933,
            "Ixx": 10.818
        }

        def ir_results_in_correct_sect_props():
            assert expected["A"] == approx(sect_props["A"], abs=0.001)
            assert expected["Ixx"] == approx(sect_props["Ixx"], abs=0.001)

        def it_results_in_correct_Pcr():
            assert signature[lengths.index(expected["xcrl"])] == approx(expected["Pcrl"], abs=0.01)
            assert signature[lengths.index(expected["xcrd"])] == approx(expected["Pcrd"], abs=0.01)

        def it_results_in_correct_signature_curve():
            assert np.allclose(expected_curve[:, 1, 0], curve[:, 0], rtol=1.e-4)

    def context_dsm_3_2_5_z_with_lips_Mx():

        @pspec_context("DSM Guide Jan 2006, Ex 3.2.2: Z-Section with Lips (Mx)")
        def describe():
            pass

        lengths, expected_curve, sect_props, signature, curve, _ = mat_file_test("zwlip_Mxr.mat")

        expected = {"xcrl": 4.1, "Mcrl": 0.85, "xcrd": 22.1, "Mcrd": 0.77, "A": 0.822, "Ixx": 7.762}

        def ir_results_in_correct_sect_props():
            assert expected["A"] == approx(sect_props["A"], abs=0.001)
            assert expected["Ixx"] == approx(sect_props["Ixx"], abs=0.001)

        def it_results_in_correct_Mcr():
            assert signature[lengths.index(expected["xcrl"])] == approx(expected["Mcrl"], abs=0.01)
            assert signature[lengths.index(expected["xcrd"])] == approx(expected["Mcrd"], abs=0.01)

        def it_results_in_correct_signature_curve():
            # TODO: Investigate why the signature curve starts mismatching at the tail end...
            assert np.allclose(expected_curve[0:-10, 1, 0], curve[0:-10, 0], rtol=1.e-3)

    def context_dsm_3_2_5_z_with_lips_P():

        @pspec_context("DSM Guide Jan 2006, Ex 3.2.2: Z-Section with Lips (P)")
        def describe():
            pass

        lengths, expected_curve, sect_props, signature, curve, _ = mat_file_test("zwlip_P.mat")

        expected = {"xcrl": 5.9, "Pcrl": 0.16, "xcrd": 18.3, "Pcrd": 0.29, "A": 0.822, "Ixx": 7.762}

        def ir_results_in_correct_sect_props():
            assert expected["A"] == approx(sect_props["A"], abs=0.001)
            assert expected["Ixx"] == approx(sect_props["Ixx"], abs=0.001)

        def it_results_in_correct_Pcr():
            assert signature[lengths.index(expected["xcrl"])] == approx(expected["Pcrl"], abs=0.01)
            assert signature[lengths.index(expected["xcrd"])] == approx(expected["Pcrd"], abs=0.01)

        def it_results_in_correct_signature_curve():
            # TODO: Investigate why the signature curve starts mismatching at the tail end...
            assert np.allclose(expected_curve[0:-10, 1, 0], curve[0:-10, 0], rtol=1.e-3)
