#!/usr/bin/env python

"""
  MechViz -- Python-based toolkit for the analysis and visualization of mechanical properties of materials

  Copyright (C) 2019-2024 by Chinedu Ekuma

  This program is free software; you can redistribute it and/or modify it under the
  terms of the GNU General Public License as published by the Free Software Foundation
  version 3 of the License.

  This program is distributed in the hope that it will be useful, but WITHOUT ANY
  WARRANTY; without even the implied warranty of MERCHANTABILITY or FITNESS FOR A
  PARTICULAR PURPOSE.  See the GNU General Public License for more details.

  E-mail: cekuma1@gmail.com

"""

import sys
import os
import numpy as np
from datetime import datetime
from os import getcwd
from find_spg import find_crystal_system
from elastic_constants import calc_elastic_constants_standalone
from ase.io import vasp
from read_parameters import readstruct,read_elastic_tensor,read_rho_dim
from stability_criteria import criteria
from elastool_elate_browser import ElateAutomation
from elastic_mechanical_properties import elastic_mechanical_properties
#import pkg_resources
from importlib.metadata import version
import warnings
warnings.filterwarnings('ignore')

from importlib.metadata import version
version = version("MechViz")
#---------------------------------------------------------
#Print out citation
def print_boxed_message(ec_file=None):
    header_footer = "+" + "-" * 78 + "+"
    spacer = "| " + " " * 76 + " |"

    # List of lines to be printed
    lines = [
        (" * CITATIONS *", True),
        ("MechViz is a standalone postprocessing toolkit of ElasTool code", False),
        ("If you have used MechViz in your research, PLEASE cite:", False),
        ("", False),  # Space after the above line
        ("ElasTool: An automated toolkit for elastic constants calculation, ", False),
        ("Z.-L. Liu, C.E. Ekuma, W.-Q. Li, J.-Q. Yang, and X.-J. Li, ", False),
        ("Computer Physics Communications 270, 108180, (2022)", False),
        ("", False),

        ("", False),  # Blank line for separation
        ("Efficient prediction of temperature-dependent elastic and", False),
        ("mechanical properties of 2D materials, S.M. Kastuar, C.E. Ekuma, Z-L. Liu,", False),
        ("Nature Scientific Report 12, 3776 (2022)", False)
    ]

    def output_line(line):
        if ec_file:
            ec_file.write(line + "\n")
        else:
            print(line)

    output_line(header_footer)
    
    for line, underline in lines:
        centered_line = line.center(76)
        output_line("| " + centered_line + " |")
        
        if underline:
            underline_str = "-" * len(centered_line)
            output_line("| " + underline_str.center(76) + " |")

    # Print footer of the box
    output_line(header_footer)



def write_line(ec_file, content, padding=1, border_char="|", filler_char=" "):
    content_width = int(max_width) - (2 * int(padding)) - 2  # Subtract 2 for the border characters
    content = content[:content_width]  # Ensure content doesn't exceed the width
    line = border_char + filler_char*padding + content.ljust(content_width) + filler_char*padding + border_char
    ec_file.write(line + "\n")





def print_banner(ec_file, version):
    # Get current date and time
    current_time = datetime.now().strftime('%H:%M:%S')
    current_date = datetime.now().strftime('%Y-%m-%d')
    conclusion_msg = f"Calculations ended at {current_time} on {current_date}"

    # Concatenate the message with the version info
    message = f"SUMMARY OF RESULTS\nusing\nMechViz Version: {version}\n{conclusion_msg}"

    # Now use the write_line function
    write_line(ec_file, '❤' * (max_width - 2), padding=0, border_char='❤', filler_char='❤')  # This will print a line of hearts
    for line in message.split('\n'):
        centered_line = line.center(max_width - 4)  # Subtract 4 for the two border characters and spaces at each end
        write_line(ec_file, centered_line, padding=1, border_char='❤')
    write_line(ec_file, '❤' * (max_width - 2), padding=0, border_char='❤', filler_char='❤')  # This will print another line of hearts

#---------------------------------------------------------


cwd = getcwd()
tubestrain_type = "Nanotube" 
pos = readstruct()


# Initialize flags
plotly_flag = False
elate_flag = False
plot_flag = False

# Iterate over all arguments (excluding the script name itself)
for arg in sys.argv[1:]:
    arg = arg.lower()
    if arg == "-plotly":
        plotly_flag = True
    elif arg == "-noplotly":
        plotly_flag = False
    elif arg == "-elate":
        elate_flag = True
    elif arg == "-plot":
        plot_flag = True

plotparameters = [plot_flag, plotly_flag]


# Now you can use plotly_flag, elate_flag, and plot_flag in your code

if elate_flag:
    browser = input("Choose a browser; when done press Ctrl+C: (chrome, firefox, edge, safari): ").lower()
    elate_instance = ElateAutomation(browser_name=browser)
    elate_instance.run()
    sys.exit(0)


rho, dimensional = read_rho_dim("massdensity_dim.dat") #You can read the rho from file or automatically obtain it from the structure
elastic_tensor = read_elastic_tensor("elastic_tensor.dat")
latt_system = find_crystal_system(pos, dimensional,tubestrain_type,False)


if dimensional == '1D':
    if latt_system == 'Nanotube':
         tubestrain_type = 'Nanotube'
    else:
        print('Choose a type of strain for the 1D system!!!')
        sys.exit(1)
  
time_start = datetime.now()

print("")
print("Reading elastic tensor parameters and structure information ...")
print("")




        
        
# Estimate the max width based on the longest expected line:
max_width = len("|WARNING: This is an empirical approx; validity needs to be checked !! |")

calc = True
if calc:

    elastic_constants_dict = calc_elastic_constants_standalone(latt_system,dimensional, elastic_tensor,{})
   
    elastic_constants_dict,hardness_values = elastic_mechanical_properties(elastic_constants_dict, pos,elastic_tensor, dimensional, latt_system,plotparameters)



    eigenvals,_ = np.linalg.eig(elastic_tensor)
    eigenvals = sorted(eigenvals)
    longdash = '-' * 55


    with open('mechviz.out', 'w') as ec_file:
        #write_line(ec_file, "")
        print_banner(ec_file,version)
        write_line(ec_file, "=" * max_width, border_char="+", filler_char="-")

        if dimensional == '1D':
            description = "            This is a %2s %s" % (
                dimensional,   tubestrain_type + ' lattice.')
        else:
            description = "            This is a %2s %s" % (
                dimensional, latt_system + ' lattice.')

        write_line(ec_file, description)
      
        write_line(ec_file, "--" * max_width, border_char="+", filler_char="-")

        print_anisotropy = False
        print_hardness = False

        try:
            if dimensional == '2D':
                Cs = np.linalg.inv(elastic_tensor)
                C11 = elastic_tensor[0,0]
                C22 = elastic_tensor[1,1]
                C12 = elastic_tensor[0,1]
                C66 = elastic_tensor[2,2]

                S11 = Cs[0, 0] 
                S22 = Cs[1, 1]
                S12 = Cs[0, 1]
                S66 = Cs[2, 2]

                B_R = 1./(S11 + S22 +2*S12)
                B_V = (C11+C22+2*C12)/4.  #Extreme Mechanics Letters, 34, 100615 (2020)
                G_R = 2./(S11 + S22 - 2*S12 + S66)
                G_V = (C11 + C22 -2*C12 +4 *C66)/8. 
                A_U = 2 * G_V / G_R + B_V / B_R - 3
                print_anisotropy = True
                print_hardness = True
            elif dimensional == '3D':
                G_V = elastic_constants_dict['G_v']
                G_R = elastic_constants_dict['G_r']
                B_V = elastic_constants_dict['B_v']
                B_R = elastic_constants_dict['B_r']
                A_U = 5 * G_V / G_R + B_V / B_R - 6
                A_C = (G_V - G_R) / (G_V + G_R)
                print_anisotropy = True
                print_hardness = True
            #else:
            #    print("Invalid dimensionality specified.")
        except:
            pass
            

        has_print_ec = False
        has_print_moduli = False
        has_print_sound = False

        for key in elastic_constants_dict.keys():
            if dimensional == '3D':
                if key[0] == 'c' and not has_print_ec:
                    write_line(ec_file, longdash)
                    write_line(ec_file, "       Elastic Constants and Mechanical Properties ")
                    write_line(ec_file, "--" * max_width, border_char="+", filler_char="-")
                    has_print_ec = True

                content_mapping = {
                    'c': "%s = %s GPa",
                    'B': "%s = %s GPa",
                    'G': "%s = %s GPa",
                    'E': "Young's modulus (%s) = %s GPa",
                    'v': "Poisson's ratio (%s) = %s",
                    'V': "Sound velocity (%s) = %s Km/s",
                    'P': "Pugh's modulus ratio (%s) = %s",
                    'M': "Lame's parameter (%s) = %s N/m",
                    'Q': "Kleinman’s parameter (%s) = %s",
                    'T': "Debye temperature (%s) = %s K",
                    'K': "Min thermal conductivity (%s) = %s W/(mK)",
                    'C': "Linear compressibility (%s) = %.2e GPa^-1", 
                    'D': "Ductility test (%s) = %s" 
 
                }

                content = content_mapping.get(key[0], None)
                if content:
                    if key[0] == "C":
                        write_line(ec_file, content % (key.capitalize(), elastic_constants_dict[key]))
                    elif key[0] not in ["D"]:
                        write_line(ec_file, content %
                                   (key.capitalize(), "%.2f" % elastic_constants_dict[key]))
                    else:
                        write_line(ec_file, content % (key.capitalize(), elastic_constants_dict[key]))  

            elif dimensional == '2D':
                print_hardness = True
                content_mapping = {
                    'c': "%s = %s N/m",
                    'Y': "Young's modulus (%s) = %s N/m",
                    'v': "Poisson's ratio (%s) = %s",
                    'B': "Stiffness constant (%s) = %s N/m",
                    'G': "Shear modulus (%s) = %s N/m",
                    'V': "Sound velocity (%s) = %s Km/s",
                    'P': "Pugh's modulus ratio (%s) = %s",
                    'L': "Layer modulus (%s) = %s N/m",
                    'T': "Debye temperature (%s) = %s K",
                    'M': "Lame's parameter (%s) = %s N/m",
                    'Q': "Kleinman’s parameter (%s) = %s",
                    'R': "Resonance frequency (%s) = %s GHz",
                    'K': "Min thermal conductivity (%s) = %s W/(mK)",
                    'C': "Linear compressibility (%s) = %.2e m/N", 
                    'D': "Ductility test (%s) = %s" 
                }

                content = content_mapping.get(key[0], None)
                if content:
                    if key[0] == "C":
                        write_line(ec_file, content % (key.capitalize(), elastic_constants_dict[key]))
                    elif key[0] not in ["D"]:
                        write_line(ec_file, content %
                                   (key.capitalize(), "%.2f" % elastic_constants_dict[key]))

                    else:
                        write_line(ec_file, content % (key.capitalize(), elastic_constants_dict[key]))  





            elif dimensional == '1D':
                content_mapping = {
                    'c': "%s = %s GPa",
                    'Y': "Young's modulus (%s) = %s GPa",
                    'v': "Poisson's ratio (%s) = %s",
                    'B': "Bulk modulus (%s) = %s GPa",
                    'G': "Shear modulus (%s) = %s GPa",
                    'C': "Compliance (%s) = %s GPa^-1",
                    'Rf': "Resonance frequency (%s) = %s GHz",
                    'V': "Sound velocity (%s) = %s Km/s",
                    'P': "Pugh's modulus ratio (%s) = %s",
                    'T': "Debye temperature (%s) = %s K",
                    'K': "Min thermal conductivity (%s) = %s W/(mK)",
                    'D': "Ductility test (%s) = %s" 
                }

                content = content_mapping.get(key[0], None)
                if content:
                    if key[0] not in ["D"]:
                        write_line(ec_file, content %
                                   (key.capitalize(), "%.2f" % elastic_constants_dict[key]))
                    else:
                        write_line(ec_file, content % (key.capitalize(), elastic_constants_dict[key]))  

       # if dimensional == '2D':
       #     energy_density_line = "Strain Energy Density = {:.3e} J/m²".format(SEDF_values)
       # else:
       #     energy_density_line = "Strain Energy Density = {:.3e} J/m³ ".format(SEDF_values)

 
        #write_line(ec_file, strain_energy_line)
        #write_line(ec_file, energy_density_line)

        if print_anisotropy:
            write_line(ec_file, longdash)
            write_line(ec_file, "Elastic anisotropy:")
            write_line(ec_file, "A_U = %s" % "%.4f" % A_U)
            if dimensional == '3D':
                write_line(ec_file, "A_C = %s" % "%.4f" % A_C)


        #print("%9.3f %9.3f %9.3f %9.3f %9.3f %9.3f" % tuple(eigenvals))



        eigen_stable = True
        if eigenvals[0] <= 0:
            #print('Eigenvalue matrix is not definite positive, crystal is mechanically unstable<br/>')
            eigen_stable = False 
        
        stable = criteria(elastic_constants_dict, latt_system)
        write_line(ec_file, "--" * max_width, border_char="+", filler_char="-")
        write_line(ec_file, "                 Structural Stability Analysis")
        write_line(ec_file, "--" * max_width, border_char="+", filler_char="-")

        lambda_headers = " ".join(["   λ_{}".format(i+1) for i in range(len(eigenvals))])
        total_width_eigenvalues = 6 * len(eigenvals) + (len(eigenvals) - 1) * 1  
        leading_space_for_eigenvalues = (total_width_eigenvalues - len(lambda_headers)) // 2


        leading_space_for_lambda = 10
        adjusted_leading_space_for_lambda = leading_space_for_lambda + leading_space_for_eigenvalues

        eigen_format = " ".join(["%6.3f" for _ in eigenvals])

        if stable: # and eigen_stable:
            write_line(ec_file, " " * adjusted_leading_space_for_lambda + lambda_headers)
            write_line(ec_file, "Eigenvalues: " + eigen_format % tuple(eigenvals))
            write_line(ec_file, "This structure is mechanically STABLE.")
        else:
            write_line(ec_file, " " * adjusted_leading_space_for_lambda + lambda_headers)
            write_line(ec_file, "Eigenvalues: " + eigen_format % tuple(eigenvals))
            write_line(ec_file, "This structure is NOT mechanically STABLE.")
            


        if print_hardness:
            if dimensional == '3D':
                H1a, H1b, H1c,H2, H3, H4, H5, H6, H7, F1, F2, F3 = hardness_values  # Unpacking the results
                write_line(ec_file, "--" * max_width, border_char="+", filler_char="-")
                write_line(ec_file, "Hardness (H) and Fracture Toughness (F) Analysis")
                write_line(ec_file, "WARNING: An empirical approximation; check validity!")
                write_line(ec_file, "--" * max_width, border_char="+", filler_char="-")
                # Printing the hardness values
                hardness_lines = [
                   "Hardness (H1a) = {:.2f} GPa;  Ref.[1]".format(H1a),
                   "Hardness (H1b) = {:.2f} GPa;  Ref.[1]".format(H1b),
                   "Hardness (H1c) = {:.2f} GPa;  Ref.[2]".format(H1c),
                   "Hardness (H2)  = {:.2f} GPa;  Ref.[3]".format(H2),
                   "Hardness (H3)  = {:.2f} GPa;  Ref.[4]".format(H3),
                   "Hardness (H4)  = {:.2f} GPa;  Ref.[1]".format(H4),
                   "Hardness (H5)  = {:.2f} GPa;  Ref.[5]".format(H5),
                   "Hardness (H6)  = {:.2f} GPa;  Ref.[6]".format(H6),
                   "Hardness (H7)  = {:.2f} GPa;  Ref.[7]".format(H7),
                   "Fracture Toughness (F1)  = {:.2f} MPa m¹/₂;  Ref.[5]".format(F1*1e3),
                   "Fracture Toughness (F2)  = {:.2f} MPa m¹/₂;  Ref.[6]".format(F2*1e3),
                   "Fracture Toughness (F3)  = {:.2f} MPa m¹/₂;  Ref.[6]".format(F3*1e3)
                ]

                for line in hardness_lines:
                   write_line(ec_file, line)

                column_widths = {
                    'Type': max(len("S"), len("I"), len("M")),
                    'Cubic': max(len("All,F1-2"), len("All,F1-2"), len("H1a,H7,F3")),
                    'Hexagonal': len("All,F1-2"),
                    'Orthorhombic': len("H2,H6,H7,F1-2"),
                    'Rhombohedral': len("All,F1-2"),
                    'General': len("H2,H6,H7,F1-2")
                }

                # Format headers
                header = "{:<{Type}}  {:<{Cubic}}  {:<{Hexagonal}}  {:<{Orthorhombic}}  {:<{Rhombohedral}}  {:<{General}}".format(
                    "", "Cubic", "Hexagonal", "Orthorhombic", "Rhombohedral", "General", **column_widths)
                divider = "--" * max_width

                recommendation_model_lines = [
                divider,
                header,
                divider,
                "{:<{Type}}  {:<{Cubic}}  {:<{Hexagonal}}  {:<{Orthorhombic}}  {:<{Rhombohedral}}  {:<{General}}".format(
                "I", "All,F1-2", "All,F1-2", "H2,H6,H7,F1-2", "All,F1-2", "H2,H6,H7,F1-2", **column_widths),
                "{:<{Type}}  {:<{Cubic}}  {:<{Hexagonal}}  {:<{Orthorhombic}}  {:<{Rhombohedral}}  {:<{General}}".format(
                "S", "All,F1-2", "All,F1-2", "H2,H6,H7,F1-2", "All,F1-2", "H5,H6,H7,F1-2", **column_widths),
                "{:<{Type}}  {:<{Cubic}}  {:<{Hexagonal}}  {:<{Orthorhombic}}  {:<{Rhombohedral}}  {:<{General}}".format(
                "M", "H1a,H7,F3", "H4,H7,F3", "H4,H7,F3", "H4,H7,F3", "H4,H7,F3", **column_widths),
                 divider
                ]

                for line in recommendation_model_lines:
                    write_line(ec_file,line)

                gap_lines = [
                    "Insulator (I)     : bandgap > 2 eV",
                    "Semiconductor (S) : bandgap < 2 eV",
                    "Metal (M)         : bandgap = 0"
                ]

                for line in gap_lines:
                    write_line(ec_file, line)

                write_line(ec_file, "--" * max_width)
                write_line(ec_file, "References")
                write_line(ec_file, "--" * max_width, border_char="+", filler_char="-")

                # Adding the references
                references = [
                     "[1] Comp. Mater. Sci. 50 (2011)",
                     "[2] Scientific Reports, 3776 (2022)",
                     "[3] MRS Bull. 23, 22 (1998)",
                     "[4] J. Phys.: Condens. Matter 22 315503 (2010)",
                     "[5] Intermetallics 19, 1275 (2011)",
                     "[6] J. Appl. Phys. 125, 065105 (2019)",
                     "[7] J. Appl. Phys. 126, 125109 (2019)"
                ]

                for ref in references:
                    write_line(ec_file, ref)


            elif dimensional == '2D':
                H1a, H1b, H1c,H2, H3, H4, H5, H6, H7, F1, F2, F3 = hardness_values 
                write_line(ec_file, "--" * max_width, border_char="+", filler_char="-")
                write_line(ec_file, "Hardness (H) and Fracture Toughness (F) Analysis")
                write_line(ec_file, "WARNING: An empirical approximation; check validity!")
                write_line(ec_file, "--" * max_width, border_char="+", filler_char="-")
                # Printing the hardness values
                hardness_lines = [
                   "Hardness (H1a) = {:.2f} N/m;  Ref.[1]".format(H1a),
                   "Hardness (H1b) = {:.2f} N/m;  Ref.[1]".format(H1b),
                   "Hardness (H1c) = {:.2f} N/m;  Ref.[2]".format(H1c),
                   "Hardness (H2)  = {:.2f} N/m;  Ref.[3]".format(H2),
                   "Hardness (H3)  = {:.2f} N/m;  Ref.[4]".format(H3),
                   "Hardness (H4)  = {:.2f} N/m;  Ref.[1]".format(H4),
                   "Hardness (H5)  = {:.2f} N/m;  Ref.[5]".format(H5),
                   "Hardness (H6)  = {:.2f} N/m;  Ref.[6]".format(H6),
                   "Hardness (H7)  = {:.2f} N/m;  Ref.[7]".format(H7),
                   "Fracture Toughness (F1)  = {:.2f} Nm⁻¹/₂; Ref.[5]".format(F1),
                   "Fracture Toughness (F2)  = {:.2f} Nm⁻¹/₂;  Ref.[6]".format(F2),
                   "Fracture Toughness (F3)  = {:.2f} Nm⁻¹/₂;  Ref.[6]".format(F3)
                ]

                for line in hardness_lines:
                   write_line(ec_file, line)

                column_widths = {
                    'Type': max(len("S"), len("I"), len("M")),
                    'Isotropy': max(len("All,F1-2"), len("H5-7,F1-2"), len("H1a,H7,F3")),
                    'Tetragonal': len("All,F1-2"),
                    'Orthotropy': len("All,F1-2"),
                    'Anisotropy': len("X-H2-4,F1-2"),
                    'General': len("H1a,H1b,H7,F1-2")
                }

                # Format headers
                header = "{:<{Type}}  {:<{Isotropy}}  {:<{Tetragonal}}  {:<{Orthotropy}}  {:<{Anisotropy}}  {:<{General}}".format(
                    "", "Isotropy", "Tetragonal", "Orthotropy", "Anisotropy", "General", **column_widths)
                divider = "--" * max_width

                recommendation_model_lines = [
                divider,
                header,
                divider,
                "{:<{Type}}  {:<{Isotropy}}  {:<{Tetragonal}}  {:<{Orthotropy}}  {:<{Anisotropy}}  {:<{General}}".format(
                "I", "All,F1-2", "All,F1-2", "All,F1-2", "X-H2-4,F1-2", "X-H2-4,F1-2", **column_widths),
                "{:<{Type}}  {:<{Isotropy}}  {:<{Tetragonal}}  {:<{Orthotropy}}  {:<{Anisotropy}}  {:<{General}}".format(
                "S", "All,F1-2", "All,F1-2", "All,F1-2", "X-H2-4,F1,F2", "H5-7,F1-2", **column_widths),
                "{:<{Type}}  {:<{Isotropy}}  {:<{Tetragonal}}  {:<{Orthotropy}}  {:<{Anisotropy}}  {:<{General}}".format(
                "M", "H1a,H7,F3", "H4,H7,F3", "H1a,H7,F3", "H1a,H7,F3", "X-H2-4,H7,F3", **column_widths),
                 divider
                ]

                for line in recommendation_model_lines:
                    write_line(ec_file,line)

                gap_lines = [
                    "Insulator (I)     : bandgap > 2 eV",
                    "Semiconductor (S) : bandgap < 2 eV",
                    "Metal (M)         : bandgap = 0"
                ]

                for line in gap_lines:
                    write_line(ec_file, line)

                write_line(ec_file, "--" * max_width)
                write_line(ec_file, "References")
                write_line(ec_file, "--" * max_width, border_char="+", filler_char="-")

                # Adding the references
                references = [
                     "[1] Comp. Mater. Sci. 50 (2011)",
                     "[2] Scientific Reports, 3776 (2022)",
                     "[3] MRS Bull. 23, 22 (1998)",
                     "[4] J. Phys.: Condens. Matter 22 315503 (2010)",
                     "[5] Intermetallics 19, 1275 (2011)",
                     "[6] J. Appl. Phys. 125, 065105 (2019)",
                     "[7] J. Appl. Phys. 126, 125109 (2019)"
                ]

                for ref in references:
                    write_line(ec_file, ref)


        write_line(ec_file, "")
        write_line(ec_file, "=" * max_width, border_char="+", filler_char="-")

        print_boxed_message(ec_file)
        ec_file.write("\n")


time_now = datetime.now()
time_used = (time_now - time_start).seconds

with open('time_used.log', 'w') as time_record_file:
    time_record_file.write("The stress calculations used %d seconds.\n" % time_used)


for line in open('mechviz.out', 'r'):
    l = line.strip('\n')
    print(l)
print("")
print("Results are also saved in the mechviz.out file.")
print("")
print("")
#print_boxed_message()
print("Well done! GOOD LUCK!")
print("")




