'''
This module provides utilities for importing and structuring hydrogen bond data generated by cpptraj.

It parses `hbond ... out <file> series` output tables, extracts residue–residue hydrogen-bond
pairs from the header, and loads the corresponding time series into adjacency-matrix
representations. These system-level matrices can then be used directly for downstream
analyses (i.e., clustering, dimensionality reduction, visualization).

Rather than functioning as a generic wrapper, the focus here is on:

* Mapping atom-level cpptraj outputs to residue-level indices.
* Constructing per-frame residue×residue H-bond matrices.
* Providing template arrays for efficient modification and integration with the broader mdsa_tools pipeline.

See Also
--------
mdsa_tools.Viz.visualize_reduction
    Plot PCA/UMAP embeddings of reduced systems.
mdsa_tools.Data_gen_hbond.create_system_representations
    Build system-level adjacency matrices from trajectories.
numpy.linalg.svd
    Linear algebra backend for certain reduction methods.

'''

import numpy as np
import mdtraj as md
from typing import Tuple, Dict

class cpptraj_hbond_import():
    '''Lightweight loader for cpptraj hydrogen-bond (series) output.

    Init takes the filepath to the desired data file and the associated topology.
    After construction, the header is parsed to determine residue–residue index
    pairs and the time-series data are loaded into memory.

    Parameters
    ----------
    filepath : str or pathlib.Path
        Path to a cpptraj `hbond ... out <file> series` output table.
    topology : str or pathlib.Path
        Path to a topology file readable by MDTraj (i.e., AMBER `prmtop`).

    Attributes
    ----------
    indices : list of tuple of int
        Parsed residue index pairs `(res1, res2)` (1-based, AMBER `resSeq` style)
        in the same order as data columns (excluding `#Frame`).
    data : np.ndarray of int, shape=(n_frames, n_pairs)
        Binary/indicator series where each column corresponds to a header pair
        in `indices`. Rows are frames.
    topology : mdtraj.Topology
        Loaded MDTraj topology associated with the series table.

    Returns
    -------
    None
        Object is initialized with parsed attributes.

    Examples
    --------
    >>> loader = cpptraj_hbond_import("hbonds.dat", "system.prmtop")
    >>> loader.indices[:3]
    [(12, 34), (12, 35), (25, 30)]
    >>> loader.data.shape
    (n_frames, n_pairs)

    Notes
    -----
    This class assumes a header where the first field is `#Frame` followed by
    columns named like `<prefix>_<res1>@<atom1>_<res2>@<atom2>`.
    '''
    def __init__(self,filepath,topology):
        
        self.indices=self.extract_headers(filepath)
        self.data=np.loadtxt(filepath, skiprows=1, usecols=range(1, len(self.indices)+1), dtype=int)
        self.topology = md.load_topology(topology) 

        return
     
    def extract_headers(self,filepath):
        '''Parse the cpptraj hydrogen-bond header to get residue–residue pairs.

        This reads only the first line of a cpptraj `hbond ... out <file> series`
        table and extracts the residue indices for each H-bond column. It expects
        a leading `#Frame` column followed by columns named like
        `<prefix>_<res1>@<atom1>_<res2>@<atom2>` (i.e., `HB_12@N_34@O`).
        The returned pairs are 1-based residue indices (AMBER `resSeq` style),
        ordered exactly as the data columns appear in the file.

        Parameters
        ----------
        filepath : str or pathlib.Path
            Path to the cpptraj `hbond` series output file. Must contain a header
            line beginning with `#Frame` and column names formatted as described
            above.

        Returns
        -------
        indices : list of tuple of int
            A list of `(res1, res2)` residue index pairs (1-based) corresponding
            to the non-`#Frame` columns in the header, in column order. These
            indices are intended to be used later to place column values into a
            residue×residue adjacency matrix at positions `[res1-1, res2-1]`.

        Raises
        ------
        FileNotFoundError
            If `filepath` does not exist.
        ValueError
            If a header token cannot be parsed into integer residue indices.

        Notes
        -----
        * Only the first line is inspected; data lines are not parsed here.
        * Column names must contain at least three underscore-separated tokens:
          a freeform prefix, `<res1>@<atom1>`, and `<res2>@<atom2>`.

        Examples
        --------
        Suppose the header line looks like::

            #Frame HB_12@N_34@O HB_12@N_35@O HB_25@O_30@H

        Then::

            indices = obj.extract_headers("hbonds.dat")
            # indices == [(12, 34), (12, 35), (25, 30)]
        '''
        filepath = filepath if filepath is not None else None

        lines=[]
        indices=[]
        with open(filepath,'r') as infile:
            for line in infile:
                lines.append(line.split())

        for col_header in lines[0]:
            if col_header !='#Frame':
                res1 = col_header.split('_')[1].split('@')[0]
                res2 = col_header.split('_')[2].split('@')[0]
                indices.append((int(res1),int(res2)))
                
        return indices

    def create_cpptraj_attributes(self,data,topology,granularity=None):
            '''Create a template array for per-frame residue adjacency matrices.

            Parameters
            ----------
            data : np.ndarray, shape=(n_frames, n_pairs)
                The loaded cpptraj series table values (without the `#Frame` column).
            topology : str or mdtraj.Topology
                Path to a topology file, or an existing `mdtraj.Topology`. If `None`,
                uses `self.topology`.
            granularity : {'residue'}, optional
                Level of coarse-graining for the template. Only 'residue' is currently
                supported; other values are reserved for future use.

            Returns
            -------
            template_array : np.ndarray, shape=(n_frames, n_res+1, n_res+1)
                A zero-initialized array containing one residue×residue matrix per
                frame. The first row and first column (index 0) store 1-based residue
                indices for convenience; the submatrix `[1:, 1:]` is the numeric
                adjacency that will be filled by `create_systems_rep`.

            Examples
            --------
            >>> template = obj.create_cpptraj_attributes(obj.data, obj.topology)
            >>> template.shape
            (n_frames, n_res+1, n_res+1)

            Notes
            -----
            We materialize a single reusable data structure (the "template") and
            fill it later for efficiency. Diagonals are initialized to 0; indexing
            labels are stored in row/column 0.
            '''

            granularity = granularity if granularity is not None else 'residue'

            #Make atom to residue dictionary 

            #Create adjacency matrix, set first row and column as residue indices, and multiply to match the number of frames
            
            topology = md.load_topology(topology) if topology is not None else self.topology

            if granularity == 'residue':

                indexes=[residue.resSeq+1 for residue in topology.residues]
                empty_array = np.zeros(shape=(len(indexes)+1,len(indexes)+1)) 

                empty_array[0,1:]=indexes
                empty_array[1:,0]=indexes

                template_array=np.repeat(empty_array[np.newaxis,:, :], data.shape[0], axis=0)

                return template_array
            
    def create_systems_rep(self,data=None,topology=None,indices=None):
        '''Fill a residue×residue H-bond matrix series from cpptraj columns.

        Parameters
        ----------
        data : np.ndarray, shape=(n_frames, n_pairs), optional
            If provided, use this series matrix; otherwise uses `self.data`.
        topology : mdtraj.Topology or str, optional
            Topology or path to topology; if `None`, uses `self.topology`.
        indices : list of tuple of int, optional
            Residue index pairs `(res1, res2)` (1-based) in column order;
            if `None`, uses `self.indices`.

        Returns
        -------
        systems : np.ndarray, shape=(n_frames, n_res+1, n_res+1)
            A time series of residue-level adjacency matrices. Row/column 0 hold
            1-based residue indices; the numeric adjacency is in `[1:, 1:]`.

        Notes
        -----
        * Self-contacts (where `res1 == res2`) are skipped and left as zeros.
        * Each column in `data` is placed at `[res1-1, res2-1]` for all frames.
        * This function does not symmetrize; if directionality matters for your
          definition, you may post-process the result.

        Examples
        --------
        >>> systems = obj.create_systems_rep()
        >>> systems.shape
        (n_frames, n_res+1, n_res+1)
        '''
        topology = topology if topology is not None else self.topology
        data = data if data is not None else self.data
        indices = indices if indices is not None else self.indices

        template_array=self.create_cpptraj_attributes(data,topology)

        iterator=0

        for col in data.T: #simply transpose so we are going column wise instead
            current_pair=indices[iterator]

            if current_pair[0]!=current_pair[1]:
                template_array[:,current_pair[0]-1,current_pair[1]-1]=col

            iterator+=1
        
        return template_array
    

if __name__ == '__main__':

        
    from mdsa_tools.Data_gen_hbond import TrajectoryProcessor as tp
    import numpy as np
    import os
    from mdsa_tools.Convenience import unrestrained_residues

    topology = '../PDBs/5JUP_N2_GCU_nowat.prmtop'
    traj = '../PDBs/CCU_GCU_10frames.mdcrd' 

    test_trajectory = tp(trajectory_path=traj,topology_path=topology)


    print("succesfully loaded current PDB from test")

    os._exit(0)

    test_atomic_system=test_trajectory.create_system_representations(test_trajectory.trajectory,granularity='atom')
    print(test_atomic_system.shape)

    test_atomic_system_no_indexes=test_atomic_system[0,1:,1:]
    print(test_atomic_system_no_indexes[test_atomic_system_no_indexes!=0])

    print('test running just the datagen file')


    #########################################
    #In house test with our own trajectories#
    #########################################

    #load in and test trajectory
    system_one_topology = '../PDBs/5JUP_N2_CGU_nowat.prmtop'
    system_one_trajectory = './CCUGCU_G34_full.mdcrd' 
    system_two_topology = '../PDBs/5JUP_N2_GCU_nowat.prmtop'
    system_two_trajectory = './CCUCGU_G34_full.mdcrd' 

    print('run one')
    test_trajectory_one = tp(trajectory_path=system_one_trajectory,topology_path=system_one_topology)
    print("tp made")
    test_system_one_ = test_trajectory_one.create_filtered_representations(residues_to_keep=unrestrained_residues)
    print("systems made")
    np.save('./full_sampling_GCU',test_system_one_)
    print(" made")


    del test_trajectory_one, test_system_one_ 

    test_trajectory_two = tp(trajectory_path=system_two_trajectory,topology_path=system_two_topology)
    print("tp made")
    test_system_two_ = test_trajectory_two.create_filtered_representations(residues_to_keep=unrestrained_residues)
    print("systems made")
    np.save('./full_sampling_CGU',test_system_two_)
    print(" made")


    #now that its loaded in try to make object
    print("intializing creation made")
    print("finished creation")

    del test_trajectory_two, test_system_two_ 
