# https://akaszynski.github.io/pyansys/ansys_control.html#initial-setup-and-example
# pip install pyansys --upgrade
import pyansys
import numpy as np
import os
import time
from scipy import sparse
from scipy.sparse import linalg


def sim_rotor(materials, angles, outer_diameter=0.250, inner_diameter=0.050, thickness=0.005, omega=1):
    layers = len(angles)
    layer_thickness = thickness / layers

    path = os.getcwd() + "\\ws"
    mapdl = pyansys.launch_mapdl(run_location=path, interactive_plotting=True, loglevel='ERROR', override=True)
    # """
    # clear existing geometry
    mapdl.finish()
    mapdl.clear()

    th = 50  # [°C]
    tn = 50  # [°C]

    mapdl.prep7()
    # define Material
    for i in range(0, len(materials)):
        mapdl.mptemp(i + 1, th)

        mapdl.mpdata("EX", i + 1, "", materials[i]["e11"])
        mapdl.mpdata("EY", i + 1, "", materials[i]["e22"])
        mapdl.mpdata("EZ", i + 1, "", materials[i]["e33"])

        mapdl.mpdata("PRXY", i + 1, "", materials[i]["pr12"])
        mapdl.mpdata("PRXZ", i + 1, "", materials[i]["pr13"])
        mapdl.mpdata("PRYZ", i + 1, "", materials[i]["pr23"])

        mapdl.mpdata("DENS", i + 1, "", materials[i]["dens"])

        mapdl.mpdata("GXY", i + 1, "", materials[i]["g12"])
        mapdl.mpdata("GXZ", i + 1, "", materials[i]["g13"])
        mapdl.mpdata("GYZ", i + 1, "", materials[i]["g23"])

        # Waermeausdenungsk
        mapdl.mpdata("ALPX", i + 1, "", materials[i]["alp11"])
        mapdl.mpdata("ALPY", i + 1, "", materials[i]["alp22"])
        mapdl.mpdata("ALPZ", i + 1, "", materials[i]["alp33"])

        # Waermeleitfaehigk
        mapdl.mpdata("KXX", i + 1, "", materials[i]["k11"])
        mapdl.mpdata("KYY", i + 1, "", materials[i]["k22"])
        mapdl.mpdata("KZZ", i + 1, "", materials[i]["k33"])

        # Quellkoeffizient
        mapdl.mpdata("BETX", i + 1, "", materials[i]["bet11"])
        mapdl.mpdata("BETY", i + 1, "", materials[i]["bet22"])
        mapdl.mpdata("BETZ", i + 1, "", materials[i]["bet33"])

        # Diffusionskoeffizient
        mapdl.mpdata("DXX", i + 1, "", materials[i]["d11"])
        mapdl.mpdata("DYY", i + 1, "", materials[i]["d22"])
        mapdl.mpdata("DZZ", i + 1, "", materials[i]["d33"])

    mapdl.et(1, "SOLID186")

    mapdl.csys(1)
    mapdl.k(1, str(inner_diameter/2), 0, str(-thickness/2))
    mapdl.k(2, str(inner_diameter/2), 0, str(thickness/2))
    mapdl.k(3, str(outer_diameter/2), 0, str(thickness/2))
    mapdl.k(4, str(outer_diameter/2), 0, str(-thickness/2))
    mapdl.l(1, 2)
    mapdl.l(2, 3)
    mapdl.l(3, 4)
    mapdl.l(4, 1)
    mapdl.al(1, 2, 3, 4)
    mapdl.k(5, 0, 0, str(thickness/2))
    mapdl.k(6, 0, 0, str(-thickness/2))
    mapdl.run("VROT,1,,,,,,5,6")

    arcdiv = 20*4
    darc = 360/arcdiv

    mapdl.lsel("S", "LOC", "X", str(inner_diameter/2))
    mapdl.lsel("R", "LOC", "Z", str(-thickness/2))
    mapdl.lesize("ALL", "", "", str(arcdiv/4))
    mapdl.lsel("S", "LOC", "X", str(inner_diameter/2))
    mapdl.lsel("R", "LOC", "Z", str(thickness/2))
    mapdl.lesize("ALL", "", "", str(arcdiv/4))
    mapdl.lsel("S", "LOC", "X", str(outer_diameter/2))
    mapdl.lsel("R", "LOC", "Z", str(-thickness/2))
    mapdl.lesize("ALL", "", "", str(arcdiv/4))
    mapdl.lsel("S", "LOC", "X", str(outer_diameter/2))
    mapdl.lsel("R", "LOC", "Z", str(thickness/2))
    mapdl.lesize("ALL", "", "", str(arcdiv/4))
    for i in range(0, 4):
        mapdl.lsel("S", "LOC", "Y", str(90*i))
        mapdl.lsel("R", "LOC", "Z", str(-thickness/2))
        mapdl.lesize("ALL", "", "", str(arcdiv/4))
        mapdl.lsel("S", "LOC", "Y", str(90*i))
        mapdl.lsel("R", "LOC", "Z", str(thickness/2))
        mapdl.lesize("ALL", "", "", str(arcdiv/4))
        mapdl.lsel("S", "LOC", "Y", str(90*i))
        mapdl.lesize("ALL", "", "", str(len(angles)))

    mapdl.lsel("ALL")
    mapdl.mshape(0, "3D")
    mapdl.mshkey(1)
    mapdl.vmesh("ALL")

    # mapdl.eplot("ALL")

    for j in range(1, len(angles)+1):
        for i in range(0, 20*4):
            mapdl.csys(1)
            mapdl.esel("S", "CENT", "Y", str(darc*(i + 0.5)))
            mapdl.esel("R", "CENT", "Z", str(thickness*(-1/2+1/len(angles)*(j-1))), str(thickness*(-1/2+1/len(angles)*j)), str(thickness/len(angles)/2))
            mapdl.local(str(j*100+i), 0, 0, 0, 0, str(angles[j-1]+darc*(i + 0.5)))
            mapdl.emodif("all", "ESYS", str(j*100+i))
            if len(angles) == len(materials):
                mapdl.emodif("all", "MAT", str(j))

    mapdl.esel("ALL")

    # Read Stress
    mapdl.run("/SOLU")
    mapdl.csys(1)
    mapdl.omega("", "", omega)
    mapdl.nsel("S", "LOC", "X", inner_diameter / 2)
    mapdl.d("ALL", "UX", 0)
    mapdl.d("ALL", "UY", 0)
    # D,ALL,ROTX,0
    # D,ALL,ROTY,0
    # D,ALL,ROTZ,0
    mapdl.nsel("R", "LOC", "Z", 0)
    mapdl.d("ALL", "UZ", 0)
    mapdl.run("ALLS")
    mapdl.outres("ALL", "ALL")
    mapdl.solve()
    mapdl.finish()

    mapdl.run("/Post1")

    result = mapdl.result
    # result.plot_nodal_solution(0, 'z', show_displacement=True)
    # result.plot_principal_nodal_stress(0, 'S1', show_edges=True, show_axes=True)

    # """
    nodenump, stress = result.principal_nodal_stress(0)
    s1max = np.nanmax(stress[:, -5])
    s2max = np.nanmax(stress[:, -4])

    nodenum, stress = result.nodal_stress(0)
    srmax = np.nanmax(stress[:, -6])
    stmax = np.nanmax(stress[:, -5])

    # Modual-Analysis
    # """
    mapdl.run("/SOLU")
    mapdl.csys(1)
    mapdl.omega("", "", omega)
    mapdl.nsel("S", "LOC", "X", inner_diameter / 2)
    mapdl.d("ALL", "UX", 0)
    mapdl.d("ALL", "UY", 0)
    mapdl.nsel("R", "LOC", "Z", 0)
    mapdl.d("ALL", "UZ", 0)
    mapdl.run("ALLS")
    mapdl.outres("ALL", "ALL")

    mapdl.antype("MODAL")
    mapdl.modopt("lanb", 3)
    mapdl.mxpand("", "", "", "yes")

    mapdl.solve()
    mapdl.finish()
    
    mapdl.run("/Post1")

    result = mapdl.result
    f = result.time_values

    mapdl.exit()
    while not mapdl._exited:
        time.sleep(0.5)

    return s1max, s2max, srmax, stmax, f[0], f[1], f[2]


"""
print(sim_rotor(
    [{
        "e11":      42.5E9  ,
        "e22":      11E9    ,
        "e33":      11E9    ,
        "pr12":     0.28    ,
        "pr13":     0.28    ,
        "pr23":     0.28    ,
        "dens":     1950    ,
        "g12":      4.2E9   ,
        "g13":      4.2E9   ,
        "g23":      2.56E9  ,
        "alp11":    5.7E-6  ,
        "alp22":    45E-6   ,
        "alp33":    45E-6   ,
        "k11":      0.72E3  ,
        "k22":      0.5E3   ,
        "k33":      0.5E3   ,
        "bet11":    0E-3,
        "bet22":    4E-3,
        "bet33":    4E-3,
        "d11":      4.4E-3,
        "d22":      3.1E-3,
        "d33":      3.1E-3
    }], [90, 90, 90, 0, 90, 90]))
# """
