'''
            SimpNMR

        Copyright (C) 2025

This program is free software: you can redistribute it and/or modify
it under the terms of the GNU Affero General Public License as published
by the Free Software Foundation, either version 3 of the License, or
(at your option) any later version.

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 Affero General Public License for more details.

You should have received a copy of the GNU Affero General Public License
along with this program.  If not, see <https://www.gnu.org/licenses/>.

This script reads temperature-dependent molar magnetic susceptibility
tensors from an ORCA output file, computes key magnetic properties
(isotropic susceptibility, axiality, rhombicity, and their ratios),
and exports the results to a structured CSV report.
'''

import argparse
import numpy as np
import numpy.linalg as la
import csv
from simpnmr import ut

def compute_matrix_properties(susc_file, section, temperature):
    '''
    Computes and prints various properties of a given matrix:
    
    - Normalizes the matrix by dividing by Temperature in K
    - Computes the traceless matrix
    - Diagonalizes the traceless matrix to obtain eigenvalues and eigenvectors
    - Sorts eigenvalues in ascending absolute value
    - Computes axiality and rhombicity
    - Prints results including eigenvectors, diagonal matrix, isotropy, axiality, rhombicity, and their ratio
    '''
    # Convert temperature argument to temperature value
    T = float(temperature)
 
    tensors = {}
    try:
        with open(susc_file, 'r') as f:
            for line in f:
                # Find the start of the QDPT section corresponding to the chosen method
                if f'QDPT WITH {section.upper()}' in line:
                    # Skip ahead until we reach the susceptibility tensor header
                    while 'TEMPERATURE DEPENDENT MOLAR MAGNETIC SUSCEPTIBILITY TENSOR' not in line:
                        line = next(f)
                    # Skip the header lines before the data block
                    for _ in range(6):
                        line = next(f)
                    # Loop over each temperature entry in the tensor data block
                    while 'TEMPERATURE/K' in line:
                        _temp = float(line.split('TEMPERATURE/K:')[1])
                        line = next(f)
                        line = next(f)
                        # Read tensor rows
                        row_1 = [float(val) for val in line.split()]
                        line = next(f)
                        row_2 = [float(val) for val in line.split()]
                        line = next(f)
                        row_3 = [float(val) for val in line.split()]
                        tensors[_temp] = np.array([row_1, row_2, row_3])
                        line = next(f)
                        line = next(f)
    except (IOError, StopIteration) as e:
        print(f"Error reading '{susc_file}': {e}")

    if not tensors:
        print(f"No data found for section '{section}' in file '{susc_file}'")

    # Select matrix for the given temperature
    try:
        matrix_K = tensors[T]
        # Record input temperature and original tensor
        orig_input_temp = T
        orig_tensor = matrix_K.copy()
    except KeyError:
        print(f"No tensor data for temperature {T} K in file '{susc_file}'")
        
    # Normalize the matrix
    matrix_norm = matrix_K / T

    # Create traceless matrix by subtracting the isotropic part from the diagonal
    iso = np.trace(matrix_norm) / 3.0
    matrix_traceless = matrix_norm - np.eye(3) * iso

    # Diagonalize the traceless matrix to compute eigenvalues and eigenvectors
    eigenvalues, eigenvectors = la.eig(matrix_traceless)

    # Sort eigenvalues by ascending absolute value
    indices = np.argsort(np.abs(eigenvalues))

    # Reorder the eigenvalues and eigenvectors according to these indices
    eigenvalues_sorted = eigenvalues[indices]
    eigenvectors_sorted = eigenvectors[:, indices]

    # Construct the diagonal matrix from the sorted eigenvalues
    D_sorted = np.diag(eigenvalues_sorted.real)

    # Calculate Axiality
    axiality = 1.5 * D_sorted[2, 2]

    # Calculate Rhombicity
    rhombicity = (D_sorted[0, 0] - D_sorted[1, 1]) / 2

    # Calculate Rhombicity/Axiality
    rhombicity_div_axiality = rhombicity / axiality

    # Export results to CSV
    output_file = f"susceptibility_{int(T)}K.csv"
    with open(output_file, 'w', newline='') as csvfile:
        writer = csv.writer(csvfile)
        # Write code generation notice
        writer.writerow(['Report generated by simpNMR'])
        writer.writerow([])
        # Write input temperature and original tensor
        writer.writerow(['Input temperature (K)', orig_input_temp])
        writer.writerow([])
        writer.writerow(['Original tensor (cm^3*K/mol)'])
        for row in orig_tensor:
            writer.writerow(row.tolist())
        writer.writerow([])
        writer.writerow(['Property', 'Value'])
        writer.writerow(['Isotropy (cm^3/mol)', round(iso, 5)])
        writer.writerow(['Axiality (cm^3/mol)', round(axiality, 5)])
        writer.writerow(['Rhombicity (cm^3/mol)', round(rhombicity, 5)])
        writer.writerow(['Rhombicity/Axiality (cm^3/mol)', round(rhombicity_div_axiality, 5)])
        writer.writerow([])
        writer.writerow(['Eigenvectors (|χ_xx| < |χ_yy| < |χ_zz|)'])
        for vec in eigenvectors_sorted.T:
            writer.writerow([round(val.real, 5) for val in vec])
        writer.writerow([])
        writer.writerow(['Diagonal matrix (|χ_xx| < |χ_yy| < |χ_zz|)'])
        for row in D_sorted:
            writer.writerow([round(val, 5) for val in row])
    print(f"Results exported to {output_file}")

def main():
    # Define command-line interface for input file and section choice
    parser = argparse.ArgumentParser(
        description=(
            'Plot χ(T) components from ORCA output\n\n'
            'Example:\n'
            '  get_susc filename.out nevpt2 298'
        ),
        formatter_class=argparse.RawDescriptionHelpFormatter
    )
    parser.add_argument(
        'susc_file',
        help='ORCA output file with susceptibility data'
    )
    parser.add_argument(
        'section',
        help='Computational method identifier for QDPT (e.g., nevpt2)'
    )
    parser.add_argument(
        'temperature',
        help='Temperature in K'
    )
    args = parser.parse_args()
    compute_matrix_properties(args.susc_file, args.section, args.temperature)

if __name__ == '__main__':
    main()