# -*- coding: utf-8 -*-

from os.path import join

import matplotlib.pyplot as plt
from numpy import array, pi, zeros

from ....Classes.Frame import Frame
from ....Classes.LamSlotWind import LamSlotWind
from ....Classes.LamSquirrelCage import LamSquirrelCage
from ....Classes.MachineDFIM import MachineDFIM
from ....Classes.Shaft import Shaft
from ....Classes.VentilationCirc import VentilationCirc
from ....Classes.VentilationPolar import VentilationPolar
from ....Classes.VentilationTrap import VentilationTrap
from ....Classes.Winding import Winding
from ....Classes.WindingUD import WindingUD
from ....Classes.WindingCW2LT import WindingCW2LT
from ....Classes.WindingDW2L import WindingDW2L
from ....Classes.MatMagnetics import MatMagnetics
from ....Classes.SlotW21 import SlotW21

from ....Tests import save_plot_path as save_path
from ....Tests.Plot.LamWind import wind_mat


"""unittest for Lamination with winding plot"""


def test_Lam_Wind_21_wind_22():
    """Test machine plot with Slot 21 and winding rad=2, tan=2
    """
    print("\nTest plot Slot 21")
    plt.close("all")
    test_obj = MachineDFIM()

    test_obj.rotor = LamSlotWind(
        Rint=0.2,
        Rext=0.5,
        is_internal=True,
        is_stator=False,
        L1=0.85,
        Nrvd=3,
        Wrvd=0.05,
    )
    test_obj.rotor.slot = SlotW21(
        Zs=6, W0=40e-3, W1=60e-3, W2=40e-3, H0=20e-3, H1=0, H2=130e-3, H1_is_rad=False,
    )
    test_obj.rotor.axial_vent.append(
        VentilationTrap(Zh=6, Alpha0=pi / 6, W1=30e-3, W2=60e-3, D0=0.05, H0=0.3)
    )
    test_obj.rotor.axial_vent.append(
        VentilationTrap(Zh=6, Alpha0=pi / 6, W1=60e-3, W2=90e-3, D0=0.05, H0=0.4)
    )
    test_obj.rotor.winding = WindingUD(user_wind_mat=wind_mat, qs=4, p=4, Lewout=60e-3)
    test_obj.rotor.mat_type.mag = MatMagnetics(Wlam=0.5e-3)
    test_obj.shaft = Shaft(Drsh=test_obj.rotor.Rint * 2, Lshaft=1)

    test_obj.stator = LamSlotWind(
        Rint=0.51,
        Rext=0.8,
        is_internal=False,
        is_stator=True,
        L1=0.85,
        Nrvd=3,
        Wrvd=0.05,
    )
    test_obj.stator.slot = SlotW21(
        Zs=18,
        W0=40e-3,
        W1=60e-3,
        W2=90e-3,
        H0=15e-3,
        H1=35e-3,
        H2=140e-3,
        H1_is_rad=False,
    )
    test_obj.stator.winding = WindingDW2L(qs=3, p=3)
    test_obj.stator.axial_vent.append(
        VentilationCirc(Zh=12, Alpha0=pi / 6, D0=50e-3, H0=0.75)
    )
    test_obj.stator.mat_type.mag = MatMagnetics(Wlam=0.5e-3)
    test_obj.stator.winding.Lewout = 60e-3
    test_obj.frame = Frame(Rint=0.8, Rext=1, Lfra=1)

    test_obj.plot()
    fig = plt.gcf()
    fig.savefig(join(save_path, "test_Lam_Wind_s21_1-Machine.png"))
    # Rotor + Stator + 2 for frame + 1 for shaft
    assert len(fig.axes[0].patches) == 91

    test_obj.rotor.plot()
    fig = plt.gcf()
    fig.savefig(join(save_path, "test_Lam_Wind_s21_2-Rotor.png"))
    # 2 for lam + 4*Zs for wind + 6 vents + 6 vents
    assert len(fig.axes[0].patches) == 38

    test_obj.stator.plot()
    fig = plt.gcf()
    fig.savefig(join(save_path, "test_Lam_Wind_s21_3-Stator.png"))
    # 2 for lam + Zs*2 for wind + 12 vents
    assert len(fig.axes[0].patches) == 50

    tooth = test_obj.rotor.slot.get_surface_tooth()
    tooth.plot(color="r")
    fig = plt.gcf()
    fig.savefig(join(save_path, "test_Lam_Wind_s21_Tooth_in.png"))

    tooth = test_obj.stator.slot.get_surface_tooth()
    tooth.plot(color="r")
    fig = plt.gcf()
    fig.savefig(join(save_path, "test_Lam_Wind_s21_Tooth_out.png"))
