# -*- coding: utf-8 -*-
"""
plot3d_geom.py
"""

import numpy as np
import matplotlib.pyplot as plt

def plot(Title, Ngeom, Ne, E_posm, E_posp, E_ifeed, E_iload):
    # figure
    strfig = 'OpenMOM - geometry (3D)'
    fig = plt.figure(strfig, figsize=(5, 5))  # 5 inches
    ax = fig.add_subplot(projection='3d')

    # marker
    marker = 'o'
    markersize = 2

    # plot
    for ie in range(Ne):
        # line
        x = [E_posm[ie, 0], E_posp[ie, 0]]
        y = [E_posm[ie, 1], E_posp[ie, 1]]
        z = [E_posm[ie, 2], E_posp[ie, 2]]

        # color and linewidth
        color = 'k'
        lw = 1
        if   E_ifeed[ie] > 0:
            lw = 2
            color = 'r'
        elif E_iload[ie] > 0:
            lw = 2
            color = 'g'

        ax.plot(x, y, z, color=color, linewidth=lw, marker=marker, markersize=markersize)

    ax.set_aspect('equal')
    ax.view_init(elev = 30, azim = 30, roll = 0)

    # label
    ax.set_xlabel('X[m]')
    ax.set_ylabel('Y[m]')
    ax.set_zlabel('Z[m]')

    # title
    ax.set_title('%s\nNo. of geometries = %d\nNo. of elements = %d' % (Title, Ngeom, Ne))

    # limit
    gmin = np.zeros(3, float)
    gmax = np.zeros(3, float)
    for m in range(3):
        gmin[m] = min(np.min(E_posm[:, m]), np.min(E_posp[:, m]))
        gmax[m] = max(np.max(E_posm[:, m]), np.max(E_posp[:, m]))
    margin = 0.05 * sum(gmax - gmin)
    gmin -= margin
    gmax += margin
    ax.set_xlim([gmin[0], gmax[0]])
    ax.set_ylim([gmin[1], gmax[1]])
    ax.set_zlim([gmin[2], gmax[2]])

    #plt.show()
