import pandas as pd
from Bio import SeqIO
import wgdi.base as base


class ancestral_karyotype:
    def __init__(self, options):
        self.mark = 'aak'
        
        # Set attributes from options
        for k, v in options:
            setattr(self, str(k), v)
            print(f"{k} = {v}")

    def run(self):
        # Load and filter data
        gff = base.newgff(self.gff)
        ancestor = base.read_classification(self.ancestor)
        gff = gff[gff['chr'].isin(ancestor[0].values.tolist())]

        # Create new gff copy and initialize required variables
        newgff = gff.copy()
        data, num = [], 1

        # Create dictionary mapping chromosome to order
        chr_arr = ancestor[3].drop_duplicates().to_list()
        chr_dict = {chr: idx + 1 for idx, chr in enumerate(chr_arr)}
        ancestor['order'] = ancestor[3].map(chr_dict)

        dict1, dict2 = {}, {}

        # Process ancestor and gff information
        for (cla, order), group in ancestor.groupby([4, 'order'], sort=[False, False]):
            for index, row in group.iterrows():
                index1 = gff[(gff['chr'] == row[0]) & (gff['order'] >= row[1]) & (gff['order'] <= row[2])].index
                newgff.loc[index1, 'chr'] = str(num)
                
                # Store results in data
                for k in index1:
                    data.append(newgff.loc[k, :].values.tolist() + [k])

            dict1[str(num)] = cla
            dict2[str(num)] = group[3].values[0]
            num += 1

        # Create dataframe from the data collected
        df = pd.DataFrame(data)

        # Filter based on peptide file
        pep = SeqIO.to_dict(SeqIO.parse(self.pep_file, "fasta"))
        df = df[df[6].isin(pep.keys())]

        # Assign new names and order
        for name, group in df.groupby(0):
            df.loc[group.index, 'order'] = range(1, len(group) + 1)
            df.loc[group.index, 'newname'] = [f"{self.mark}{name}g{i:05d}" for i in range(1, len(group) + 1)]

        # Set data types and sort
        df['order'] = df['order'].astype(int)
        df = df[[0, 'newname', 1, 2, 3, 'order', 6]].sort_values(by=[0, 'order'])

        # Save output files
        df.to_csv(self.ancestor_gff, sep="\t", index=False, header=None)
        lens = df.groupby(0).max()[[2, 'order']]
        lens.to_csv(self.ancestor_lens, sep="\t", header=None)

        # Add extra columns and save final results
        lens[1] = 1
        lens['color'] = lens.index.map(dict2)
        lens['class'] = lens.index.map(dict1)
        lens[[1, 'order', 'color', 'class']].to_csv(self.ancestor_file, sep="\t", header=None)

        # Update peptide sequences with new IDs and save
        id_dict = df.set_index(6).to_dict()['newname']
        seqs = []

        for seq_record in SeqIO.parse(self.pep_file, "fasta"):
            if seq_record.id in id_dict:
                seq_record.id = id_dict[seq_record.id]
                seqs.append(seq_record)

        SeqIO.write(seqs, self.ancestor_pep, "fasta")
