import mdtraj as md
import numpy as np
import os
import pandas as pd

def concatenate_trajectories(base_dir, protein_name, n_segments):
    protein_name = protein_name + '_0'
    gt_trajs = []
    gen_trajs = []
    
    
    for i in range(n_segments):
        segment_dir = f'{base_dir}/{i}'

        gt_traj = md.load(f'{segment_dir}/{protein_name}_gt.xtc', 
                         top=f'{segment_dir}/{protein_name}_gt.pdb')
        gt_trajs.append(gt_traj)
        

        gen_traj = md.load(f'{segment_dir}/{protein_name}.xtc',
                          top=f'{segment_dir}/{protein_name}.pdb')
        gen_trajs.append(gen_traj)
    

    concatenated_gt = md.join(gt_trajs)
    concatenated_gen = md.join(gen_trajs)
    

    temp_dir = f'{base_dir}/gen'
    os.makedirs(temp_dir, exist_ok=True)
    

    concatenated_gt.save(f'{temp_dir}/{protein_name}_gt.xtc')
    concatenated_gen.save(f'{temp_dir}/{protein_name}.xtc')
    

    gt_trajs[0][0].save(f'{temp_dir}/{protein_name}_gt.pdb')
    gen_trajs[0][0].save(f'{temp_dir}/{protein_name}.pdb')
    
    return temp_dir

def main():

    base_dir = 'generated xtc folder'

    df = pd.read_csv('/mdcath/test.csv')
    protein_names = df['name'].tolist()
    
    n_segments = 20 
    
    for protein_name in protein_names:
        
        print(f"\nProcessing {protein_name}...")
        
        try:
            temp_dir = concatenate_trajectories(base_dir, protein_name, n_segments)
            print(f"Successfully concatenated trajectories for {protein_name}")
            print(f"Files saved in {temp_dir}")
        except Exception as e:
            print(f"Error processing {protein_name}: {str(e)}")

if __name__ == "__main__":
    main()