import gurobipy as gp
from gurobipy import GRB

from prettify import red, green, warn, blue

GRANTHAM_MATRIX = { ('A', 'A'): '0',
  ('A', 'C'): '195',
  ('A', 'D'): '126',
  ('A', 'E'): '107',
  ('A', 'F'): '113',
  ('A', 'G'): '60',
  ('A', 'H'): '86',
  ('A', 'I'): '94',
  ('A', 'K'): '106',
  ('A', 'L'): '96',
  ('A', 'M'): '84',
  ('A', 'N'): '111',
  ('A', 'P'): '27',
  ('A', 'Q'): '91',
  ('A', 'R'): '112',
  ('A', 'S'): '99',
  ('A', 'T'): '58',
  ('A', 'V'): '64',
  ('A', 'W'): '148',
  ('A', 'Y'): '112',
  ('C', 'A'): '195',
  ('C', 'C'): '0',
  ('C', 'D'): '154',
  ('C', 'E'): '170',
  ('C', 'F'): '205',
  ('C', 'G'): '159',
  ('C', 'H'): '174',
  ('C', 'I'): '198',
  ('C', 'K'): '202',
  ('C', 'L'): '198',
  ('C', 'M'): '196',
  ('C', 'N'): '139',
  ('C', 'P'): '169',
  ('C', 'Q'): '154',
  ('C', 'R'): '180',
  ('C', 'S'): '112',
  ('C', 'T'): '149',
  ('C', 'V'): '192',
  ('C', 'W'): '215',
  ('C', 'Y'): '194',
  ('D', 'A'): '126',
  ('D', 'C'): '154',
  ('D', 'D'): '0',
  ('D', 'E'): '45',
  ('D', 'F'): '177',
  ('D', 'G'): '94',
  ('D', 'H'): '81',
  ('D', 'I'): '168',
  ('D', 'K'): '101',
  ('D', 'L'): '172',
  ('D', 'M'): '160',
  ('D', 'N'): '23',
  ('D', 'P'): '108',
  ('D', 'Q'): '61',
  ('D', 'R'): '96',
  ('D', 'S'): '65',
  ('D', 'T'): '85',
  ('D', 'V'): '152',
  ('D', 'W'): '181',
  ('D', 'Y'): '160',
  ('E', 'A'): '107',
  ('E', 'C'): '170',
  ('E', 'D'): '45',
  ('E', 'E'): '0',
  ('E', 'F'): '140',
  ('E', 'G'): '98',
  ('E', 'H'): '40',
  ('E', 'I'): '134',
  ('E', 'K'): '56',
  ('E', 'L'): '138',
  ('E', 'M'): '126',
  ('E', 'N'): '42',
  ('E', 'P'): '93',
  ('E', 'Q'): '29',
  ('E', 'R'): '54',
  ('E', 'S'): '80',
  ('E', 'T'): '65',
  ('E', 'V'): '121',
  ('E', 'W'): '152',
  ('E', 'Y'): '122',
  ('F', 'A'): '113',
  ('F', 'C'): '205',
  ('F', 'D'): '177',
  ('F', 'E'): '140',
  ('F', 'F'): '0',
  ('F', 'G'): '153',
  ('F', 'H'): '100',
  ('F', 'I'): '21',
  ('F', 'K'): '102',
  ('F', 'L'): '22',
  ('F', 'M'): '28',
  ('F', 'N'): '158',
  ('F', 'P'): '114',
  ('F', 'Q'): '116',
  ('F', 'R'): '97',
  ('F', 'S'): '155',
  ('F', 'T'): '103',
  ('F', 'V'): '50',
  ('F', 'W'): '40',
  ('F', 'Y'): '22',
  ('G', 'A'): '60',
  ('G', 'C'): '159',
  ('G', 'D'): '94',
  ('G', 'E'): '98',
  ('G', 'F'): '153',
  ('G', 'G'): '0',
  ('G', 'H'): '98',
  ('G', 'I'): '135',
  ('G', 'K'): '127',
  ('G', 'L'): '138',
  ('G', 'M'): '127',
  ('G', 'N'): '80',
  ('G', 'P'): '42',
  ('G', 'Q'): '87',
  ('G', 'R'): '125',
  ('G', 'S'): '56',
  ('G', 'T'): '59',
  ('G', 'V'): '109',
  ('G', 'W'): '184',
  ('G', 'Y'): '147',
  ('H', 'A'): '86',
  ('H', 'C'): '174',
  ('H', 'D'): '81',
  ('H', 'E'): '40',
  ('H', 'F'): '100',
  ('H', 'G'): '98',
  ('H', 'H'): '0',
  ('H', 'I'): '94',
  ('H', 'K'): '32',
  ('H', 'L'): '99',
  ('H', 'M'): '87',
  ('H', 'N'): '68',
  ('H', 'P'): '77',
  ('H', 'Q'): '24',
  ('H', 'R'): '29',
  ('H', 'S'): '89',
  ('H', 'T'): '47',
  ('H', 'V'): '84',
  ('H', 'W'): '115',
  ('H', 'Y'): '83',
  ('I', 'A'): '94',
  ('I', 'C'): '198',
  ('I', 'D'): '168',
  ('I', 'E'): '134',
  ('I', 'F'): '21',
  ('I', 'G'): '135',
  ('I', 'H'): '94',
  ('I', 'I'): '0',
  ('I', 'K'): '102',
  ('I', 'L'): '5',
  ('I', 'M'): '10',
  ('I', 'N'): '149',
  ('I', 'P'): '95',
  ('I', 'Q'): '109',
  ('I', 'R'): '97',
  ('I', 'S'): '142',
  ('I', 'T'): '89',
  ('I', 'V'): '29',
  ('I', 'W'): '61',
  ('I', 'Y'): '33',
  ('K', 'A'): '106',
  ('K', 'C'): '202',
  ('K', 'D'): '101',
  ('K', 'E'): '56',
  ('K', 'F'): '102',
  ('K', 'G'): '127',
  ('K', 'H'): '32',
  ('K', 'I'): '102',
  ('K', 'K'): '0',
  ('K', 'L'): '107',
  ('K', 'M'): '95',
  ('K', 'N'): '94',
  ('K', 'P'): '103',
  ('K', 'Q'): '53',
  ('K', 'R'): '26',
  ('K', 'S'): '121',
  ('K', 'T'): '78',
  ('K', 'V'): '97',
  ('K', 'W'): '110',
  ('K', 'Y'): '85',
  ('L', 'A'): '96',
  ('L', 'C'): '198',
  ('L', 'D'): '172',
  ('L', 'E'): '138',
  ('L', 'F'): '22',
  ('L', 'G'): '138',
  ('L', 'H'): '99',
  ('L', 'I'): '5',
  ('L', 'K'): '107',
  ('L', 'L'): '0',
  ('L', 'M'): '15',
  ('L', 'N'): '153',
  ('L', 'P'): '98',
  ('L', 'Q'): '113',
  ('L', 'R'): '102',
  ('L', 'S'): '145',
  ('L', 'T'): '92',
  ('L', 'V'): '32',
  ('L', 'W'): '61',
  ('L', 'Y'): '36',
  ('M', 'A'): '84',
  ('M', 'C'): '196',
  ('M', 'D'): '160',
  ('M', 'E'): '126',
  ('M', 'F'): '28',
  ('M', 'G'): '127',
  ('M', 'H'): '87',
  ('M', 'I'): '10',
  ('M', 'K'): '95',
  ('M', 'L'): '15',
  ('M', 'M'): '0',
  ('M', 'N'): '142',
  ('M', 'P'): '87',
  ('M', 'Q'): '101',
  ('M', 'R'): '91',
  ('M', 'S'): '135',
  ('M', 'T'): '81',
  ('M', 'V'): '21',
  ('M', 'W'): '67',
  ('M', 'Y'): '36',
  ('N', 'A'): '111',
  ('N', 'C'): '139',
  ('N', 'D'): '23',
  ('N', 'E'): '42',
  ('N', 'F'): '158',
  ('N', 'G'): '80',
  ('N', 'H'): '68',
  ('N', 'I'): '149',
  ('N', 'K'): '94',
  ('N', 'L'): '153',
  ('N', 'M'): '142',
  ('N', 'N'): '0',
  ('N', 'P'): '91',
  ('N', 'Q'): '46',
  ('N', 'R'): '86',
  ('N', 'S'): '46',
  ('N', 'T'): '65',
  ('N', 'V'): '133',
  ('N', 'W'): '174',
  ('N', 'Y'): '143',
  ('P', 'A'): '27',
  ('P', 'C'): '169',
  ('P', 'D'): '108',
  ('P', 'E'): '93',
  ('P', 'F'): '114',
  ('P', 'G'): '42',
  ('P', 'H'): '77',
  ('P', 'I'): '95',
  ('P', 'K'): '103',
  ('P', 'L'): '98',
  ('P', 'M'): '87',
  ('P', 'N'): '91',
  ('P', 'P'): '0',
  ('P', 'Q'): '76',
  ('P', 'R'): '103',
  ('P', 'S'): '74',
  ('P', 'T'): '38',
  ('P', 'V'): '68',
  ('P', 'W'): '147',
  ('P', 'Y'): '110',
  ('Q', 'A'): '91',
  ('Q', 'C'): '154',
  ('Q', 'D'): '61',
  ('Q', 'E'): '29',
  ('Q', 'F'): '116',
  ('Q', 'G'): '87',
  ('Q', 'H'): '24',
  ('Q', 'I'): '109',
  ('Q', 'K'): '53',
  ('Q', 'L'): '113',
  ('Q', 'M'): '101',
  ('Q', 'N'): '46',
  ('Q', 'P'): '76',
  ('Q', 'Q'): '0',
  ('Q', 'R'): '43',
  ('Q', 'S'): '68',
  ('Q', 'T'): '42',
  ('Q', 'V'): '96',
  ('Q', 'W'): '130',
  ('Q', 'Y'): '99',
  ('R', 'A'): '112',
  ('R', 'C'): '180',
  ('R', 'D'): '96',
  ('R', 'E'): '54',
  ('R', 'F'): '97',
  ('R', 'G'): '125',
  ('R', 'H'): '29',
  ('R', 'I'): '97',
  ('R', 'K'): '26',
  ('R', 'L'): '102',
  ('R', 'M'): '91',
  ('R', 'N'): '86',
  ('R', 'P'): '103',
  ('R', 'Q'): '43',
  ('R', 'R'): '0',
  ('R', 'S'): '110',
  ('R', 'T'): '71',
  ('R', 'V'): '96',
  ('R', 'W'): '101',
  ('R', 'Y'): '77',
  ('S', 'A'): '99',
  ('S', 'C'): '112',
  ('S', 'D'): '65',
  ('S', 'E'): '80',
  ('S', 'F'): '155',
  ('S', 'G'): '56',
  ('S', 'H'): '89',
  ('S', 'I'): '142',
  ('S', 'K'): '121',
  ('S', 'L'): '145',
  ('S', 'M'): '135',
  ('S', 'N'): '46',
  ('S', 'P'): '74',
  ('S', 'Q'): '68',
  ('S', 'R'): '110',
  ('S', 'S'): '0',
  ('S', 'T'): '58',
  ('S', 'V'): '124',
  ('S', 'W'): '177',
  ('S', 'Y'): '144',
  ('T', 'A'): '58',
  ('T', 'C'): '149',
  ('T', 'D'): '85',
  ('T', 'E'): '65',
  ('T', 'F'): '103',
  ('T', 'G'): '59',
  ('T', 'H'): '47',
  ('T', 'I'): '89',
  ('T', 'K'): '78',
  ('T', 'L'): '92',
  ('T', 'M'): '81',
  ('T', 'N'): '65',
  ('T', 'P'): '38',
  ('T', 'Q'): '42',
  ('T', 'R'): '71',
  ('T', 'S'): '58',
  ('T', 'T'): '0',
  ('T', 'V'): '69',
  ('T', 'W'): '128',
  ('T', 'Y'): '92',
  ('V', 'A'): '64',
  ('V', 'C'): '192',
  ('V', 'D'): '152',
  ('V', 'E'): '121',
  ('V', 'F'): '50',
  ('V', 'G'): '109',
  ('V', 'H'): '84',
  ('V', 'I'): '29',
  ('V', 'K'): '97',
  ('V', 'L'): '32',
  ('V', 'M'): '21',
  ('V', 'N'): '133',
  ('V', 'P'): '68',
  ('V', 'Q'): '96',
  ('V', 'R'): '96',
  ('V', 'S'): '124',
  ('V', 'T'): '69',
  ('V', 'V'): '0',
  ('V', 'W'): '88',
  ('V', 'Y'): '55',
  ('W', 'A'): '148',
  ('W', 'C'): '215',
  ('W', 'D'): '181',
  ('W', 'E'): '152',
  ('W', 'F'): '40',
  ('W', 'G'): '184',
  ('W', 'H'): '115',
  ('W', 'I'): '61',
  ('W', 'K'): '110',
  ('W', 'L'): '61',
  ('W', 'M'): '67',
  ('W', 'N'): '174',
  ('W', 'P'): '147',
  ('W', 'Q'): '130',
  ('W', 'R'): '101',
  ('W', 'S'): '177',
  ('W', 'T'): '128',
  ('W', 'V'): '88',
  ('W', 'W'): '0',
  ('W', 'Y'): '37',
  ('Y', 'A'): '112',
  ('Y', 'C'): '194',
  ('Y', 'D'): '160',
  ('Y', 'E'): '122',
  ('Y', 'F'): '22',
  ('Y', 'G'): '147',
  ('Y', 'H'): '83',
  ('Y', 'I'): '33',
  ('Y', 'K'): '85',
  ('Y', 'L'): '36',
  ('Y', 'M'): '36',
  ('Y', 'N'): '143',
  ('Y', 'P'): '110',
  ('Y', 'Q'): '99',
  ('Y', 'R'): '77',
  ('Y', 'S'): '144',
  ('Y', 'T'): '92',
  ('Y', 'V'): '55',
  ('Y', 'W'): '37',
  ('Y', 'Y'): '0',
  ('U', 'A'): '301',
  ('U', 'C'): '301',
  ('U', 'D'): '301',
  ('U', 'E'): '301',
  ('U', 'F'): '301',
  ('U', 'G'): '301',
  ('U', 'H'): '301',
  ('U', 'I'): '301',
  ('U', 'K'): '301',
  ('U', 'L'): '301',
  ('U', 'M'): '301',
  ('U', 'N'): '301',
  ('U', 'P'): '301',
  ('U', 'Q'): '301',
  ('U', 'R'): '301',
  ('U', 'S'): '301',
  ('U', 'T'): '301',
  ('U', 'V'): '301',
  ('U', 'W'): '301',
  ('U', 'Y'): '301',
  ('U', 'U'): '0',
  ('A', 'U'): '301',
  ('C', 'U'): '301',
  ('D', 'U'): '301',
  ('E', 'U'): '301',
  ('F', 'U'): '301',
  ('G', 'U'): '301',
  ('H', 'U'): '301',
  ('I', 'U'): '301',
  ('K', 'U'): '301',
  ('L', 'U'): '301',
  ('M', 'U'): '301',
  ('N', 'U'): '301',
  ('P', 'U'): '301',
  ('Q', 'U'): '301',
  ('R', 'U'): '301',
  ('S', 'U'): '301',
  ('T', 'U'): '301',
  ('V', 'U'): '301',
  ('W', 'U'): '301',
  ('Y', 'U'): '301',
  ('U', 'U'): '0'
}

# attach index id to each amino acid to differentiate them in the LP
def prepare(aa_set1, aa_set2):
    count_dict = {}
    new_set1 = []
    new_set2 = []
    for a in aa_set1:
        if a not in count_dict:
            count_dict[a] = 0
        else:
            count_dict[a] += 1
        new_set1.append(f"{a}{count_dict[a]}")

    for a in aa_set2:
        if a not in count_dict:
            count_dict[a] = 0
        else:
            count_dict[a] += 1
        new_set2.append(f"{a}{count_dict[a]}")

    return (new_set1, new_set2)

def min_distance_LP(aa_set1, aa_set2):
    (aa_set1, aa_set2) = prepare(aa_set1, aa_set2)
    lp = gp.Model("Min_Grantham", env=gp.Env(params={"OutputFlag":0}))
    lp.setParam("OutputFlag", 0)

    # OBJ: Min. SUM(x_i * distance(v_i, v_j)) + PENALTY*SUM(1-v_i)
    # x_i: boolean, whether or not a pair is selected. Eg. 'A' is paired with 'C'
    # v_i: boolean, whether or not amino acid v_i is paired

    xi_dict = {}
    vi_dict = {}
    for a in aa_set1:
        vi_dict[a] = lp.addVar(vtype=GRB.BINARY)
    for b in aa_set2:
        vi_dict[b] = lp.addVar(vtype=GRB.BINARY)

    for a in aa_set1:
        row_xi = []
        for b in aa_set2:
            xi_varhandle = lp.addVar(vtype=GRB.BINARY)
            xi_dict[(a,b)] = xi_varhandle
            row_xi.append(xi_varhandle)
            # x_i <= v_a, x_i <= v_b
            # if pair x_i is selected, v_a and v_b must be selected
            lp.addConstr(xi_varhandle <= vi_dict[a])
            lp.addConstr(xi_varhandle <= vi_dict[b])

    # if v_i is chosen, there must be exactly one x_i which covers it
    # v_i == sum(x_i with v_i in it)
    for a in aa_set1:
        covered = []
        for (x,y) in list(xi_dict.keys()):
            if a==x:
                covered.append(xi_dict[(x,y)])
        lp.addConstr(vi_dict[a] == sum(covered))
    
    for b in aa_set2:
        covered = []
        for (x,y) in list(xi_dict.keys()):
            if b==y:
                covered.append(xi_dict[(x,y)])
        lp.addConstr(vi_dict[b] == sum(covered))

    # set objective: Min. SUM(x_i * distance(v_i, v_j)) + PENALTY*SUM(1-v_i)
    # Penalty: pairing is ALWAYS cheaper than unpaired
    PENALTY = 300
    coefs = {}
    for (a,b) in list(xi_dict.keys()):
        aa1 = a[0]
        aa2 = b[0]
        distance = int(GRANTHAM_MATRIX[(aa1,aa2)])
        coefs[(a,b)] = distance
    #lp.setObjective(sum([coef*xi_dict[(a,b)] for ((a,b), coef) in list(coefs.items())]), GRB.MINIMIZE)
    lp.setObjective(sum([coef*xi_dict[(a,b)] for ((a,b), coef) in list(coefs.items())]) \
                   + PENALTY*sum([(1-vi) for vi in list(vi_dict.values())]), GRB.MINIMIZE)

    lp.optimize()

    if lp.status != GRB.OPTIMAL:
        print(red("LP infeasible. Quitting."))
        return (0, [], [])

    selected_pairs = []
    unpaired_aas = []
    rev_xi_dict = {v : k for k,v in list(xi_dict.items())}
    rev_vi_dict = {v : k for k,v in list(vi_dict.items())}
    for v in lp.getVars():
        if v in rev_xi_dict and v.X>0:
            selected_pairs.append(rev_xi_dict[v])
        if v in rev_vi_dict and v.X==0:
            unpaired_aas.append(rev_vi_dict[v])
    
    return (lp.objVal, selected_pairs, unpaired_aas)


# calculate the Grantham distance (approx) between two sets of amino acids
# Amino Acids sets should be in the form ['A', 'Y', 'W']
def min_distance_measure(aa_set1, aa_set2):
    aa_set1 = sorted(aa_set1)
    aa_set2 = sorted(aa_set2)

    set1_matched = []
    set2_matched = []
    # filter out amino acids which match
    for a in aa_set1:
        match_b = False
        for b in aa_set2:
            if a==b:
                set1_matched.append(a)
                set2_matched.append(b)
                match_b = b
                break
        if match_b: aa_set2.remove(match_b)
    [aa_set1.remove(x) for x in set1_matched]
    
    distance = 0
    set1_matched = []
    for a in aa_set1:
        if aa_set2:
            min_distance = int(GRANTHAM_MATRIX[(a,aa_set2[0])])
            pair_aa = aa_set2[0]
        else: break

        for b in aa_set2:
            curr_distance = int(GRANTHAM_MATRIX[(a, b)])
            if curr_distance < min_distance:
                min_distance = curr_distance
                pair_aa = b
        distance += min_distance
        aa_set2.remove(pair_aa)
    [aa_set1.remove(x) for x in set1_matched]
        
    # penalty distance when number of amino acids does not match
    PENALTY = 100
    total_penalty = abs(len(aa_set1) - len(aa_set2)) * PENALTY
    distance += total_penalty

    return distance