from dataclasses import dataclass, field
from pathlib import Path
from modeller import *
from modeller.automodel import *
import time
from typing import List
import sys
import glob
import pyfastx
from Bio import SeqIO
from Bio.Seq import Seq
from Bio.SeqRecord import SeqRecord

@dataclass
class PDBModeler:
    structure_file: Path
    fasta_file: Path
    outdir: Path
    chain: str
    num_loop: int = 2
    md_level: str = 'refine.fast' # refine.very_fast or refine.slow optional

    def __post_init__(self):
        self.structure = self.structure_file.stem
        self.sequence = self.fasta_file.stem 
        self.ali_file = self.fasta_to_ali()

    @staticmethod
    def find_non_dash_indices(seq):
        start = next((i for i, c in enumerate(seq) if c != '-'), None)
        end = next((i for i, c in enumerate(reversed(seq)) if c != '-'), None)
        if start is not None and end is not None:
            end = len(seq) - end
        return start, end
    
    @staticmethod
    def align_sequences(file: Path) -> Path:
        fx = pyfastx.Fasta(file.as_posix(), build_index=True)
        assert len(fx) == 2
        seqs = [seq for seq in fx]
        
        # 确定哪条链需要裁剪
        if seqs[0].seq.startswith('-') or seqs[0].seq.endswith('-'):
            trim_index = 0
        elif seqs[1].seq.startswith('-') or seqs[1].seq.endswith('-'):
            trim_index = 1
        else:
            # 如果两条链都不需要裁剪，就直接返回原文件
            return None

        start, end = PDBModeler.find_non_dash_indices(seqs[trim_index].seq)
        
        # 根据确定的裁剪位置裁剪两条链
        trimmed_seqs = []
        for seq in seqs:
            trimmed_seq = seq.seq[start:end]
            trimmed_seqs.append(SeqRecord(Seq(trimmed_seq), id=seq.name, description=""))

        # 选择没有'-'的序列
        selected_seq = None
        for seq_record in trimmed_seqs:
            if '-' not in seq_record.seq:
                selected_seq = seq_record
                break

        assert selected_seq is not None, "no sequence without '-' found"
        assert not selected_seq.seq.startswith('-') and not selected_seq.seq.endswith('-'), "selected sequence should not start or end with '-'"

        # Write the selected sequence to a new FASTA file using Biopython
        new_fasta_file = file.with_suffix('.selected.fasta')
        with open(new_fasta_file, 'w') as output_handle:
            SeqIO.write([selected_seq], output_handle, "fasta")

        return new_fasta_file

    def make_model(self): # 单模板建模
        print("***************************************************")
        print("md_level ====", self.md_level)
        print("***************************************************")

        time_start = time.time()

        env1 = Environ()
        mdl = Model(env1, file=self.structure_file.as_posix(), model_segment=(f'FIRST:{self.chain}', f'LAST:{self.chain}'))
        aln = Alignment(env1)
        aln.append_model(mdl, align_codes=self.structure, atom_files=self.structure_file.as_posix())
        aln.append(file=self.ali_file.as_posix(), align_codes=self.sequence)
        aln.align2d()
        aln.write(file=(self.outdir / 'alignment1.ali').as_posix(), alignment_format='PIR')
        aln.write(file=(self.outdir / 'alignment1.pap').as_posix(), alignment_format='PAP')

        # save alignment in FASTA format
        aln.write(file=(self.outdir / 'alignment1.fasta').as_posix(), alignment_format='FASTA')
        slice_fasta = PDBModeler.align_sequences(self.outdir / 'alignment1.fasta')
        if slice_fasta:
            slice_ali = self.outdir / 'alignment_slice.ali'
            fx = pyfastx.Fasta(slice_fasta.as_posix(), build_index=True)
            assert len(fx) == 1, "FASTA file should contain only one sequence"
            PDBModeler.write_ali(slice_ali, fx[0].name, fx[0].seq)
            env2 = Environ()
            mdl2 = Model(env2, file=self.structure_file.as_posix(), model_segment=(f'FIRST:{self.chain}', f'LAST:{self.chain}'))
            aln2 = Alignment(env2)
            aln2.append_model(mdl2, align_codes=self.structure, atom_files=self.structure_file.as_posix())
            aln2.append(file=slice_ali.as_posix(), align_codes=self.sequence)
            aln2.align2d()
            aln2.write(file=(self.outdir / 'alignment2.ali').as_posix(), alignment_format='PIR')
            aln2.write(file=(self.outdir / 'alignment2.pap').as_posix(), alignment_format='PAP')

        log.verbose()

        # choose ali file
        fix_ali_file = self.outdir / 'alignment2.ali' if (self.outdir / 'alignment2.ali').exists() else self.outdir / 'alignment1.ali'

        env3 = Environ()
        env3.io.atom_files_directory = ['.']
        loop_model = LoopModel(env3,
                               alnfile=fix_ali_file.as_posix(),
                               knowns=self.structure,
                               sequence=self.sequence,
                               loop_assess_methods=(assess.DOPE, assess.GA341))
        # 设置循环模型数量
        # 数量规则：(end - start) + 1
        loop_model.loop.starting_model = 1
        loop_model.loop.ending_model = self.num_loop
        # 设置 MD 优化函数为 "refine.slow" 或 "refine.fast
        if self.md_level.strip() == 'refine.slow':
            loop_model.loop.md_level = refine.slow
        elif self.md_level.strip() == 'refine.very_fast':
            loop_model.loop.md_level = refine.very_fast
        elif self.md_level.strip() == 'refine.fast':
            loop_model.loop.md_level = refine.fast

        # 调用 LoopModel 的 make 方法
        loop_model.make()
        end_time = time.time()
        print(f"Time cost: {end_time - time_start}s")

        # 获取所有成功生成的模型文件的路径
        model_files = self.get_model_files(loop_model)
        if model_files:
            print(f"Model files: {[file.name for file in model_files]}")
        else:
            print("No model files found.")

        return model_files
    
    def find_pdb95_fsa_file(self) -> Path:
        """在 Conda 环境中查找 pdb95.fsa 文件的路径。"""
        # 获取当前 Python 解释器的路径
        python_executable_path = Path(sys.executable)

        # 获取 Conda 环境的根目录
        conda_env_root = python_executable_path.parent.parent

        # 获取可能的 Modeller 目录
        modeller_dirs = list(conda_env_root.glob("lib/modeller-*/examples/commands"))

        # 选择最新版本的 Modeller 目录
        modeller_dirs.sort(reverse=True)
        if modeller_dirs:
            latest_modeller_dir = modeller_dirs[0]
            pdb95_fsa_path = latest_modeller_dir / "pdb95.fsa"
            return pdb95_fsa_path
        else:
            raise FileNotFoundError("Modeller directory not found.")

    def get_model_files(self, loop_model) -> List[Path]:
        # 检查 loop_model.loop.outputs 列表，收集所有成功生成的模型文件
        model_files = []
        for output in loop_model.loop.outputs:
            if output.get('failure') is None:
                model_files.append(Path(output.get('name')))
        return model_files
    
    @staticmethod
    def write_ali(ali_file: Path, description: str, sequence: str):
        with open(ali_file, 'w') as f:
            f.write(f'>P1;{description}\n')
            f.write(f'sequence:{description}:::::::0.00: 0.00\n')
            f.write(f'{sequence}*')

    def fasta_to_ali(self) -> Path:
        if not self.outdir.exists():
            self.outdir.mkdir(parents=True, exist_ok=True)

        ali_file = self.outdir / f'{self.sequence}.ali'
        if ali_file.exists():
            ali_file.unlink()

        fx = pyfastx.Fasta(self.fasta_file.as_posix(), build_index=True)
        assert len(fx) == 1, "FASTA file should contain only one sequence"
        PDBModeler.write_ali(ali_file, self.sequence, fx[0].seq)
        return ali_file

if __name__ == "__main__":
    import argparse
    parser = argparse.ArgumentParser(description="Build model by Modeller")
    parser.add_argument("-s", "--structure", help="Structure file")
    parser.add_argument("-o", "--outdir", help="Output directory")
    parser.add_argument("-f", "--fasta", help="Fasta file")
    parser.add_argument("-n", "--num_loop", help="Number of loop model")
    parser.add_argument("-m", "--md_level", help="MD level")
    parser.add_argument("-c", "--chain", help="Chain ID")
    args = parser.parse_args()

    modeler = PDBModeler(Path(args.structure), Path(args.fasta), Path(args.outdir),
                         args.chain, int(args.num_loop), args.md_level)
    modeler.make_model()

# test command
# python build_modellel.py -s ../5sws_fixer.pdb -o ./5swsmodellerfix -f ../rcsb_pdb_5SWS.fasta -n 1 -m refine.very_fast -c D
# python build_modeller.py -s 1j8h.pdb -o ./1j8hmodellerfix -f ./1j8h_D.fasta -n 1 -m refine.very_fast -c D