from __future__ import print_function

import copy
import numpy as np
import numpy.testing as npt
import unittest
import time
import wisdem.floatingse.loading as sP
import matplotlib.pyplot as plt
from mpl_toolkits.mplot3d import Axes3D
from wisdem.commonse.vertical_cylinder import get_nfull
from wisdem.commonse import gravity as g


NSECTIONS = 5
NHEIGHT = NSECTIONS + 1
NPTS = get_nfull(NHEIGHT)


opt = {}
opt["gamma_f"] = 1.35
opt["gamma_m"] = 1.1
opt["gamma_n"] = 1.0
opt["gamma_b"] = 1.1
opt["gamma_fatigue"] = 1.755

opt["frame3dd"] = {}
opt["frame3dd"]["shear"] = True
opt["frame3dd"]["geom"] = False
opt["frame3dd"]["dx"] = -1
opt["frame3dd"]["nM"] = 6
opt["frame3dd"]["Mmethod"] = 1
opt["frame3dd"]["lump"] = 0
opt["frame3dd"]["tol"] = 1e-7

opt["run_modal"] = False

opt["tower"] = {}
opt["tower"]["buckling_length"] = 30.0

opt["columns"] = {}
opt["columns"]["main"] = {}
opt["columns"]["offset"] = {}
opt["columns"]["main"]["buckling_length"] = opt["columns"]["offset"]["buckling_length"] = 30.0


def DrawTruss(mytruss):
    mynodes = {}
    for k in range(len(mytruss.myframe.nx)):
        mynodes[mytruss.myframe.nnode[k]] = np.r_[mytruss.myframe.nx[k], mytruss.myframe.ny[k], mytruss.myframe.nz[k]]
    myelem = []
    for k in range(len(mytruss.myframe.eN1)):
        myelem.append((mytruss.myframe.eN1[k], mytruss.myframe.eN2[k]))

    fig = plt.figure()
    ax = fig.add_subplot(111, projection="3d")
    for e in myelem:
        xs = np.array([mynodes[e[0]][0], mynodes[e[1]][0]])
        ys = np.array([mynodes[e[0]][1], mynodes[e[1]][1]])
        zs = np.array([mynodes[e[0]][2], mynodes[e[1]][2]])
        ax.plot(xs, ys, zs)
    # ax.auto_scale_xyz([-10, 10], [-10, 10], [-30, 50])
    plt.show()


def getInputs():
    inputs = {}
    discrete_inputs = {}

    inputs["main_E_full"] = inputs["offset_E_full"] = inputs["tower_E_full"] = 200e9 * np.ones(NPTS - 1)
    inputs["main_G_full"] = inputs["offset_G_full"] = inputs["tower_G_full"] = 79.3e9 * np.ones(NPTS - 1)
    inputs["main_rho_full"] = inputs["offset_rho_full"] = inputs["tower_rho_full"] = 7850.0 * np.ones(NPTS - 1)
    inputs["main_sigma_y_full"] = inputs["offset_sigma_y_full"] = inputs["tower_sigma_y_full"] = 345e6 * np.ones(
        NPTS - 1
    )

    inputs["hsig_wave"] = 10.0

    inputs["rho_water"] = 1025.0

    discrete_inputs["cross_attachment_pontoons"] = True
    discrete_inputs["lower_attachment_pontoons"] = True
    discrete_inputs["upper_attachment_pontoons"] = True
    discrete_inputs["lower_ring_pontoons"] = True
    discrete_inputs["upper_ring_pontoons"] = True
    discrete_inputs["outer_cross_pontoons"] = True

    inputs["number_of_offset_columns"] = 3
    inputs["connection_ratio_max"] = 0.25

    inputs["number_of_mooring_connections"] = 3
    inputs["mooring_lines_per_connection"] = 1
    inputs["mooring_neutral_load"] = 1e2 * np.ones((15, 3))
    inputs["mooring_stiffness"] = 1e10 * np.eye(6)
    inputs["mooring_moments_of_inertia"] = np.ones(6)

    inputs["fairlead"] = 7.5
    inputs["fairlead_radius"] = 1.0
    inputs["fairlead_support_outer_diameter"] = 2.0
    inputs["fairlead_support_wall_thickness"] = 1.0

    inputs["main_pontoon_attach_upper"] = 1.0
    inputs["main_pontoon_attach_lower"] = 0.0

    inputs["radius_to_offset_column"] = 15.0
    inputs["pontoon_outer_diameter"] = 2.0
    inputs["pontoon_wall_thickness"] = 1.0

    colz = np.r_[np.linspace(-15, -10, 7)[:-1], np.linspace(-10, 0, 4)[:-1], np.linspace(0, 10, 7)]
    inputs["main_z_full"] = colz.copy()
    inputs["main_d_full"] = 2 * 10.0 * np.ones(NPTS)
    inputs["main_t_full"] = 0.1 * np.ones(NPTS - 1)
    inputs["main_mass"] = 1e2 * np.ones(NPTS - 1)
    inputs["main_buckling_length"] = 2.0 * np.ones(NPTS - 1)
    inputs["main_displaced_volume"] = 1e2 * np.ones(NPTS - 1)
    inputs["main_hydrostatic_force"] = np.zeros(NPTS - 1)
    inputs["main_hydrostatic_force"][0] = inputs["main_displaced_volume"].sum() * g * inputs["rho_water"]
    inputs["main_center_of_buoyancy"] = -10.0
    inputs["main_center_of_mass"] = -6.0
    inputs["main_Px"] = 50.0 * np.ones(NPTS)
    inputs["main_Py"] = np.zeros(NPTS)
    inputs["main_Pz"] = np.zeros(NPTS)
    inputs["main_qdyn"] = 70.0 * np.ones(NPTS)

    colz = np.r_[np.linspace(-15, 0, 10)[:-1], np.linspace(0, 2.5, 4)[:-1], np.linspace(2.5, 10, 4)]
    inputs["offset_z_full"] = colz.copy()
    inputs["offset_d_full"] = 2 * 2.0 * np.ones(NPTS)
    inputs["offset_t_full"] = 0.05 * np.ones(NPTS - 1)
    inputs["offset_mass"] = 1e1 * np.ones(NPTS - 1)
    inputs["offset_buckling_length"] = 2.0 * np.ones(NPTS - 1)
    inputs["offset_displaced_volume"] = 1e1 * np.ones(NPTS - 1)
    inputs["offset_hydrostatic_force"] = np.zeros(NPTS - 1)
    inputs["offset_hydrostatic_force"][0] = inputs["main_displaced_volume"].sum() * g * inputs["rho_water"]
    inputs["offset_center_of_buoyancy"] = -5.0
    inputs["offset_center_of_mass"] = -3.0
    inputs["offset_Px"] = 50.0 * np.ones(NPTS)
    inputs["offset_Py"] = np.zeros(NPTS)
    inputs["offset_Pz"] = np.zeros(NPTS)
    inputs["offset_qdyn"] = 70.0 * np.ones(NPTS)

    inputs["tower_z_full"] = np.linspace(0, 90, NPTS)
    inputs["tower_d_full"] = 2 * 7.0 * np.ones(NPTS)
    inputs["tower_t_full"] = 0.5 * np.ones(NPTS - 1)
    inputs["tower_mass_section"] = 2e2 * np.ones(NPTS - 1)
    inputs["tower_buckling_length"] = 25.0
    inputs["tower_center_of_mass"] = 50.0
    inputs["tower_Px"] = 50.0 * np.ones(NPTS)
    inputs["tower_Py"] = np.zeros(NPTS)
    inputs["tower_Pz"] = np.zeros(NPTS)
    inputs["tower_qdyn"] = 70.0 * np.ones(NPTS)

    inputs["rna_force"] = 6e1 * np.ones(3)
    inputs["rna_moment"] = 7e2 * np.ones(3)
    inputs["rna_mass"] = 6e1
    inputs["rna_cg"] = np.array([3.05, 2.96, 2.13])
    inputs["rna_I"] = np.array([3.05284574e9, 2.96031642e9, 2.13639924e7, 0.0, 2.89884849e7, 0.0])

    inputs["material_cost_rate"] = 1.0
    inputs["painting_cost_rate"] = 10.0
    inputs["labor_cost_rate"] = 2.0
    return inputs, discrete_inputs


class TestFrame(unittest.TestCase):
    def setUp(self):
        self.inputs, self.discrete_inputs = getInputs()
        self.outputs = {}
        self.discrete_outputs = {}

        self.outputs["pontoon_stress"] = np.zeros(70)

        self.mytruss = sP.FloatingFrame(
            n_height_main=NHEIGHT, n_height_off=NHEIGHT, n_height_tow=NHEIGHT, modeling_options=opt
        )

    def tearDown(self):
        self.mytruss = None

    def testStandard(self):
        self.inputs["radius_to_offset_column"] = 20.0
        self.inputs["fairlead_radius"] = 30.0
        # self.inputs['offset_z_full'] = np.array([-15.0, -10.0, -5.0, 0.0, 2.5, 3.0])
        self.mytruss.compute(self.inputs, self.outputs, self.discrete_inputs, self.discrete_outputs)

        npt.assert_equal(self.outputs["main_connection_ratio"], 0.25 - 0.1)
        npt.assert_equal(self.outputs["offset_connection_ratio"], 0.25 - 0.5)
        # DrawTruss(self.mytruss)

    def testSpar(self):
        self.inputs["cross_attachment_pontoons"] = False
        self.inputs["lower_attachment_pontoons"] = False
        self.inputs["upper_attachment_pontoons"] = False
        self.inputs["lower_ring_pontoons"] = False
        self.inputs["upper_ring_pontoons"] = False
        self.inputs["outer_cross_pontoons"] = False
        self.inputs["number_of_offset_columns"] = 0

        # self.mytruss.compute(self.inputs, self.outputs, self.discrete_inputs, self.discrete_outputs)

        # DrawTruss(self.mytruss)

    def testOutputsIncremental(self):
        ncyl = self.inputs["number_of_offset_columns"]
        R_semi = self.inputs["radius_to_offset_column"]
        Rmain = 0.5 * self.inputs["main_d_full"][0]
        Roff = 0.5 * self.inputs["offset_d_full"][0]
        Ro = 0.5 * self.inputs["pontoon_outer_diameter"]
        Ri = Ro - self.inputs["pontoon_wall_thickness"]
        rho = self.inputs["main_rho_full"][0]
        rhoW = self.inputs["rho_water"]

        self.inputs["fairlead_radius"] = 0.0
        self.inputs["fairlead_outer_diameter"] = 0.0
        self.inputs["fairlead_wall_thickness"] = 0.0
        self.inputs["mooring_lines_per_column"] = 0.0
        self.inputs["number_of_mooring_connections"] = 0.0
        self.discrete_inputs["cross_attachment_pontoons"] = False
        self.discrete_inputs["lower_attachment_pontoons"] = True
        self.discrete_inputs["upper_attachment_pontoons"] = False
        self.discrete_inputs["lower_ring_pontoons"] = False
        self.discrete_inputs["upper_ring_pontoons"] = False
        self.discrete_inputs["outer_cross_pontoons"] = False
        self.inputs["pontoon_main_attach_upper"] = 1.0
        self.inputs["pontoon_main_attach_lower"] = 0.0
        self.inputs["material_cost_rate"] = 6.25
        self.inputs["painting_cost_rate"] = 0.0
        self.inputs["labor_cost_rate"] = 0.0

        self.mytruss.compute(self.inputs, self.outputs, self.discrete_inputs, self.discrete_outputs)
        V = np.pi * Ro * Ro * (R_semi - Rmain - Roff) * ncyl
        m = np.pi * (Ro * Ro - Ri * Ri) * (R_semi - Rmain - Roff) * ncyl * rho
        self.assertAlmostEqual(self.outputs["pontoon_displacement"], V)
        self.assertAlmostEqual(self.outputs["pontoon_center_of_buoyancy"], -15.0)
        self.assertAlmostEqual(self.outputs["pontoon_center_of_mass"], -15.0)
        self.assertAlmostEqual(self.outputs["pontoon_mass"], m)
        self.assertAlmostEqual(self.outputs["pontoon_cost"], m * 6.25, 2)

        self.discrete_inputs["lower_ring_pontoons"] = True
        self.mytruss.compute(self.inputs, self.outputs, self.discrete_inputs, self.discrete_outputs)

        V = np.pi * Ro * Ro * ncyl * (R_semi * (1 + np.sqrt(3)) - Rmain - 3 * Roff)
        m = np.pi * (Ro * Ro - Ri * Ri) * ncyl * rho * (R_semi * (1 + np.sqrt(3)) - Rmain - 3 * Roff)
        self.assertAlmostEqual(self.outputs["pontoon_displacement"], V)
        self.assertAlmostEqual(self.outputs["pontoon_center_of_buoyancy"], -15.0)
        self.assertAlmostEqual(self.outputs["pontoon_center_of_mass"], -15.0)
        self.assertAlmostEqual(self.outputs["pontoon_mass"], m)
        self.assertAlmostEqual(self.outputs["pontoon_cost"], m * 6.25, 2)

        self.discrete_inputs["upper_attachment_pontoons"] = True  # above waterline
        self.mytruss.compute(self.inputs, self.outputs, self.discrete_inputs, self.discrete_outputs)

        V = np.pi * Ro * Ro * ncyl * (R_semi * (1 + np.sqrt(3)) - Rmain - 3 * Roff)
        m = np.pi * (Ro * Ro - Ri * Ri) * ncyl * rho * (R_semi * (2 + np.sqrt(3)) - 2 * Rmain - 4 * Roff)
        # cg = ((-15)*(1 + np.sqrt(3)) + 10) / (2+np.sqrt(3))
        self.assertAlmostEqual(self.outputs["pontoon_displacement"], V)
        self.assertAlmostEqual(self.outputs["pontoon_center_of_buoyancy"], -15.0)
        # self.assertAlmostEqual(self.outputs['pontoon_center_of_mass'], cg)
        self.assertAlmostEqual(self.outputs["pontoon_mass"], m)
        self.assertAlmostEqual(self.outputs["pontoon_cost"], m * 6.25, 2)

        self.discrete_inputs["upper_ring_pontoons"] = True
        self.mytruss.compute(self.inputs, self.outputs, self.discrete_inputs, self.discrete_outputs)

        V = np.pi * Ro * Ro * ncyl * (R_semi * (1 + np.sqrt(3)) - Rmain - 3 * Roff)
        m = np.pi * (Ro * Ro - Ri * Ri) * ncyl * rho * 2 * (R_semi * (1 + np.sqrt(3)) - Rmain - 3 * Roff)
        self.assertAlmostEqual(self.outputs["pontoon_displacement"], V)
        self.assertAlmostEqual(self.outputs["pontoon_center_of_buoyancy"], -15.0)
        self.assertAlmostEqual(self.outputs["pontoon_center_of_mass"], -2.5)
        self.assertAlmostEqual(self.outputs["pontoon_mass"], m)
        self.assertAlmostEqual(self.outputs["pontoon_cost"], m * 6.25, 2)

        # self.inputs['cross_attachment_pontoons'] = True
        # self.mytruss.compute(self.inputs, self.outputs, self.discrete_inputs, self.discrete_outputs)

        # L = np.sqrt(R_semi*R_semi + 25*25)
        # k = 15. / 25.
        # V = np.pi * Ro*Ro * ncyl * (k*L + R_semi * (1 + np.sqrt(3)))
        # m = np.pi * (Ro*Ro-Ri*Ri) * ncyl * rho * (L + R_semi * 2 * (1 + np.sqrt(3)))
        # self.assertAlmostEqual(self.outputs['pontoon_displacement'], V)
        # self.assertAlmostEqual(self.outputs['pontoon_mass'], m)
        # self.assertAlmostEqual(self.outputs['pontoon_cost'], m*6.25, 2)

        # self.inputs['outer_cross_pontoons'] = True
        # self.mytruss.compute(self.inputs, self.outputs, self.discrete_inputs, self.discrete_outputs)

    def testForces(self):
        self.inputs["offset_z_full"] = np.linspace(-1e-2, 0.0, NPTS)
        self.inputs["main_z_full"] = np.linspace(-1e-2, 0.0, NPTS)
        self.inputs["tower_z_full"] = np.linspace(0, 1e-2, NPTS)
        self.inputs["fairlead"] = 0.0
        self.inputs["number_of_offset_columns"] = 0
        self.inputs["main_mass"] = 1.0 * np.ones(NPTS - 1)
        self.inputs["main_center_of_mass"] = 0.0
        self.inputs["main_hydrostatic_force"] = 1e-12 * np.ones(NPTS - 1)
        self.inputs["offset_mass"] = 0.0 * np.ones(NPTS - 1)
        self.inputs["offset_center_of_mass"] = 0.0
        self.inputs["offset_hydrostatic_force"] = 1e-12 * np.ones(NPTS - 1)
        self.inputs["tower_mass_section"] = 1.0 * np.ones(NPTS - 1)
        self.inputs["tower_center_of_mass"] = 0.0
        self.inputs["rna_mass"] = 1.0
        self.inputs["rna_force"] = 10.0 * np.ones(3)
        self.inputs["rna_moment"] = 20.0 * np.ones(3)
        self.inputs["rna_cg"] = np.array([0.0, 0.0, 0.0])
        self.inputs["rna_I"] = np.zeros(6)
        self.inputs["main_Px"] = 0.0 * np.ones(NPTS - 1)
        self.inputs["offset_Px"] = 0.0 * np.ones(NPTS - 1)
        self.inputs["tower_Px"] = 0.0 * np.ones(NPTS - 1)
        self.inputs["main_qdyn"] = 0.0 * np.ones(NPTS)
        self.inputs["offset_qdyn"] = 0.0 * np.ones(NPTS)
        self.inputs["tower_qdyn"] = 0.0 * np.ones(NPTS)
        self.discrete_inputs["cross_attachment_pontoons"] = False
        self.discrete_inputs["lower_attachment_pontoons"] = False
        self.discrete_inputs["upper_attachment_pontoons"] = False
        self.discrete_inputs["lower_ring_pontoons"] = False
        self.discrete_inputs["upper_ring_pontoons"] = False
        self.discrete_inputs["outer_cross_pontoons"] = False
        self.inputs["main_pontoon_attach_upper"] = 1.0
        self.inputs["main_pontoon_attach_lower"] = 0.0
        self.inputs["rho_water"] = 1e-12
        self.inputs["number_of_mooring_connections"] = 3
        self.inputs["mooring_lines_per_connection"] = 1
        self.inputs["mooring_neutral_load"] = 10.0 * np.ones((15, 3))
        self.inputs["mooring_moments_of_inertia"] = np.ones(6)
        self.inputs["fairlead_radius"] = 10.1
        self.inputs["fairlead_support_outer_diameter"] = 2 * np.sqrt(2.0 / np.pi)
        self.inputs["fairlead_support_wall_thickness"] = np.sqrt(2.0 / np.pi) - np.sqrt(1.0 / np.pi)
        self.inputs["main_rho_full"] = self.inputs["offset_rho_full"] = self.inputs["tower_rho_full"] = 20.0 * np.ones(
            NPTS - 1
        )
        self.inputs["radius_to_offset_column"] = 1.0

        goodRun = False
        kiter = 0
        while goodRun == False:
            kiter += 1
            self.mytruss.compute(self.inputs, self.outputs, self.discrete_inputs, self.discrete_outputs)
            if self.outputs["substructure_mass"] < 1e3:
                goodRun = True
            if kiter > 3:
                self.assertTrue(goodRun)

        msub = (NPTS - 1) + 3 * 20.0 * 1.0 * 0.1
        mtowrna = (NPTS - 1) + 1
        mtot = msub + mtowrna

        self.assertAlmostEqual(self.outputs["substructure_mass"], msub, 5)
        self.assertAlmostEqual(self.outputs["structural_mass"], mtot, 4)
        npt.assert_almost_equal(self.outputs["total_force"], 10 * 3 + np.array([10.0, 10.0, 10 - mtot * g]), decimal=1)
        npt.assert_almost_equal(self.outputs["total_moment"], np.array([20.0, 20.0, 20.0]), decimal=0)

        self.inputs["rna_cg"] = np.array([5.0, 5.0, 5.0])
        goodRun = False
        kiter = 0
        while goodRun == False:
            kiter += 1
            self.mytruss.compute(self.inputs, self.outputs, self.discrete_inputs, self.discrete_outputs)
            if self.outputs["substructure_mass"] < 1e3:
                goodRun = True
            if kiter > 10:
                self.assertTrue(goodRun)
        npt.assert_almost_equal(self.outputs["total_force"], 3 * 10 + np.array([10.0, 10.0, 10 - mtot * g]), decimal=1)
        # self.assertAlmostEqual(self.outputs['total_moment'][-1], 20.0)


class TestModal(unittest.TestCase):
    def testRunModal(self):
        inputs, discrete_inputs = getInputs()
        outputs = {}
        discrete_outputs = {}
        outputs["pontoon_stress"] = np.zeros(70)

        newopt = copy.copy(opt)
        newopt["run_modal"] = True
        mytruss = sP.FloatingFrame(
            n_height_main=NHEIGHT, n_height_off=NHEIGHT, n_height_tow=NHEIGHT, modeling_options=newopt
        )

        mytruss.compute(inputs, outputs, discrete_inputs, discrete_outputs)
        self.assertTrue(True)


class TestSandbox(unittest.TestCase):
    def setUp(self):
        self.inputs, self.discrete_inputs = getInputs()
        self.outputs = {}
        self.discrete_outputs = {}
        self.outputs["pontoon_stress"] = np.zeros(70)

        self.mytruss = sP.FloatingFrame(
            n_height_main=NHEIGHT, n_height_off=NHEIGHT, n_height_tow=NHEIGHT, modeling_options=opt
        )

    def tearDown(self):
        self.mytruss = None

    def testBadInput(self):
        self.inputs["number_of_offset_columns"] = 1
        self.mytruss.compute(self.inputs, self.outputs, self.discrete_inputs, self.discrete_outputs)
        self.assertEqual(self.outputs["substructure_mass"], 1e30)

        self.inputs["number_of_offset_columns"] = 2
        self.mytruss.compute(self.inputs, self.outputs, self.discrete_inputs, self.discrete_outputs)
        self.assertEqual(self.outputs["substructure_mass"], 1e30)

        self.inputs["number_of_offset_columns"] = 8
        self.mytruss.compute(self.inputs, self.outputs, self.discrete_inputs, self.discrete_outputs)
        self.assertEqual(self.outputs["substructure_mass"], 1e30)

        self.inputs["number_of_offset_columns"] = 3
        self.inputs["main_z_full"][-2] = self.inputs["main_z_full"][-3] + 1e-12
        self.mytruss.compute(self.inputs, self.outputs, self.discrete_inputs, self.discrete_outputs)
        self.assertEqual(self.outputs["substructure_mass"], 1e30)

    def testCombinations(self):
        self.inputs["radius_to_offset_column"] = 30.0

        for nc in [0, 3, 4, 5, 6, 7]:

            for cap in [True, False]:

                for lap in [True, False]:

                    for uap in [True, False]:

                        for lrp in [True, False]:

                            for urp in [True, False]:

                                for ocp in [True, False]:
                                    self.inputs["number_of_offset_columns"] = nc
                                    self.inputs["number_of_mooring_connections"] = np.maximum(nc, 3)
                                    self.inputs["mooring_lines_per_connection"] = 1
                                    self.discrete_inputs["cross_attachment_pontoons"] = cap
                                    self.discrete_inputs["lower_attachment_pontoons"] = lap
                                    self.discrete_inputs["upper_attachment_pontoons"] = uap
                                    self.discrete_inputs["lower_ring_pontoons"] = lrp
                                    self.discrete_inputs["upper_ring_pontoons"] = urp
                                    self.discrete_inputs["outer_cross_pontoons"] = ocp
                                    if (nc > 0) and (not cap) and (not lap) and (not uap):
                                        continue
                                    if (nc > 0) and (ocp) and (not lrp):
                                        continue

                                    self.mytruss.compute(
                                        self.inputs, self.outputs, self.discrete_inputs, self.discrete_outputs
                                    )
                                    if self.outputs["substructure_mass"] == 1e30:
                                        print(nc, cap, lap, uap, lrp, urp, ocp)
                                    self.assertNotEqual(self.outputs["substructure_mass"], 1e30)
                                    time.sleep(1e-4)


def suite():
    suite = unittest.TestSuite()
    suite.addTest(unittest.makeSuite(TestFrame))
    suite.addTest(unittest.makeSuite(TestModal))
    suite.addTest(unittest.makeSuite(TestSandbox))
    return suite


if __name__ == "__main__":
    unittest.TextTestRunner().run(suite())
