from Bio.PDB import PDBParser, PDBIO, Structure
from io import StringIO
import json
from env_info import get_db
import pandas as pd
import os
from Bio.PDB import PDBParser, PDBIO, Select
from io import StringIO
import random
import math

class NoAltLocSelect(Select):
    def accept_atom(self, atom):
        altloc = atom.get_altloc()
        coord = atom.get_coord()
        return (altloc in (' ', 'A')) and all(coord)

def clean_pdb(pdb_string):
    parser = PDBParser(QUIET=True)
    structure = parser.get_structure("structure", StringIO(pdb_string))

    output = StringIO()
    io = PDBIO()
    io.set_structure(structure)
    io.save(output, select=NoAltLocSelect())

    return output.getvalue()


def format_longer_pdb(pdb_data_string_B, R, t):
    parser = PDBParser(QUIET=True)
    string_io = StringIO(pdb_data_string_B)
    structure = parser.get_structure("2", string_io)
    string_io.close()
    for model in structure:
        for chain in model:
            for residue in chain:
                for atom in residue:
                    coord = atom.get_coord()  # Original coordinate as a NumPy array
                    coord = R @ coord + t  # Apply rotation and translation
                    atom.set_coord(coord)  # Update atom coordinate

    io = PDBIO()
    io.set_structure(structure)
    string_io = StringIO() # Create a string buffer
    io.save(string_io)  # Write the structure to the string buffer
    pdb_data_string_B_updated = string_io.getvalue() # Get the string from the buffer
    string_io.close()

    return pdb_data_string_B_updated


def generate_combined_pdb(pdb_data_string_A, pdb_data_string_B, R, t):
    
    parser = PDBParser(QUIET=True)
    string_io = StringIO(pdb_data_string_A)
    structure1 = parser.get_structure("1", string_io)
    string_io.close()
    string_io = StringIO(pdb_data_string_B)
    structure2 = parser.get_structure("2", string_io)
    string_io.close()
    
    combined_structure = Structure.Structure('combined')

    # Add models from both structures
    model1 = structure1[0]
    model2 = structure2[0]

    # Renumber atoms globally
    atom_serial = 1
    atom_lines = []

    pdb_A_atoms_num = 0

    for model_id, model in enumerate([model1, model2]):
        atom_lines.append(f"MODEL    {model_id}\n")
        for chain in model:
            for residue in chain:
                for atom in residue:
                    if model_id == 0:
                        line = format_atom_line(atom, atom_serial)
                        pdb_A_atoms_num += 1
                    else:
                        line = format_atom_line(atom, atom_serial, R, t)
                    atom_lines.append(line)
                    atom_serial += 1
        atom_lines.append("ENDMDL\n")

    output_buffer = StringIO()
    output_buffer.writelines(atom_lines)
    pdb_data_string = output_buffer.getvalue()
    output_buffer.close()
    
    return pdb_data_string, pdb_A_atoms_num

def load_html_template(filepath):
    with open(filepath, 'r', encoding='utf-8') as file:
        html_template = file.read()
    return html_template

def write_html(template_path, placeholders, output_filename):
    # Load the HTML template from file
    html_template = load_html_template(template_path)
    
    # Replace placeholders with actual values
    for placeholder, value in placeholders.items():
        html_template = html_template.replace(f"{{{{ {placeholder} }}}}", value)
    
    # Write the final HTML to the output file
    with open(output_filename, 'w', encoding='utf-8') as file:
        file.write(html_template)

def extract_PDB_string(pdb_name, pdb_path): 
    parser = PDBParser(QUIET=True)
    io = PDBIO()
    # Load structures
    pdb_data = parser.get_structure(pdb_name, pdb_path)
    io.set_structure(pdb_data)
    string_io = StringIO() # Create a string buffer
    io.save(string_io)  # Write the structure to the string buffer
    pdb_data_string = string_io.getvalue() # Get the string from the buffer
    string_io.close()
    
    return pdb_data_string


def single_user_write_sitecontext_result_file(distance_dict, shorter_PDB_name, longer_PDB_name, f_text,
                                              shorter_PDB_string, longer_PDB_string, correspondence_list, corr_10, corr_25, corr_50,
                                              uncorrespondence_list_B, output_folder):
    
    if not os.path.exists(output_folder):
        os.makedirs(output_folder)   
    name = list(distance_dict.keys())[0]
    df = pd.DataFrame.from_dict(distance_dict, orient="index")
    df.to_csv(os.path.join(output_folder, f'similarity_scores_using_{f_text}.csv'))
    
    with open(os.path.join(output_folder, f'correspondence_using_{f_text}.text'), 'w') as f:
        f.write(f"{shorter_PDB_name:<8},{longer_PDB_name:<8}\n")
        for atom1_id, atom2_id in correspondence_list:
            f.write(f"{atom1_id:<8},{atom2_id:<8}\n")
        f.write(f"Unmatched for {longer_PDB_name}\n")
        for atom2_id in uncorrespondence_list_B:
            f.write(f"{atom2_id:<8}\n")
    correspondence_len = len(correspondence_list)
    placeholders = {
    'feature_type': f_text, 
    'shorter_PDB_name': shorter_PDB_name,
    'longer_PDB_name': longer_PDB_name,
    'shorter_PDB_string': shorter_PDB_string, 
    'longer_PDB_string': longer_PDB_string, 
    'correspondence_content_10': json.dumps([correspondence_list[i] for i in corr_10]),
    'correspondence_content_25': json.dumps([correspondence_list[i] for i in corr_25]),
    'correspondence_content_50': json.dumps([correspondence_list[i] for i in corr_50]),
    'correspondence_content': json.dumps(correspondence_list),
    'uncorrespondence_content': json.dumps(uncorrespondence_list_B)
    }
    write_html('3Dmol_viewer/html_templates/3Dmol_viewer_ribbon_only.html', 
            placeholders, os.path.join(output_folder, f'3Dmol_viewer_{f_text}.html'))
