import os
import pandas as pd
from prescient.transforms.functional import anarci_numbering
from prescient.constants import RANGES_AHO


def map_to_aho(aho_seq_chain, indexing=0):
    #print(aho_seq_chain)
    mapping = {}
    count=indexing
    for i_aho, s in enumerate(aho_seq_chain):
        if s=='-':
            continue
        else:
            mapping[count]=i_aho
            count+=1
    return mapping


def annotate_cdrs_for_sequence(heavy, light):
    #print('seq', heavy)
    heavy_aho = anarci_numbering(heavy)
    light_aho = anarci_numbering(light)
    map_h = map_to_aho(heavy_aho[0])
    map_aho_to_h = {v:k for k,v in map_h.items()}
    map_l = map_to_aho(light_aho[0])
    map_aho_to_l = {v:k for k,v in map_l.items()}
    #print(map_aho_to_l)
    #print(RANGES_AHO)
    indices = []
    for name in ['H1', 'H2', 'H3']:
        r0, r1 = RANGES_AHO[name]
        indices += [t for t in range(r0, r1+1)]
    
    #print(map_aho_to_h)
    indices_h = [map_aho_to_h[i] for i in indices if i in map_aho_to_h]
    indices = []
    for name in ['L1', 'L2', 'L3']:
        r0, r1 = RANGES_AHO[name]
        indices += [t for t in range(r0, r1+1)]

    indices_l = [map_aho_to_l[i] for i in indices  if i in map_aho_to_l]
    return indices_h, indices_l
    