#!/usr/bin/env python

import sys
import re
import os
import traceback
from lxml import etree
from Chain import Chain
from Entity import Entity
from mmcif_helper import *

def get_refmac(root):
    return root.xpath("//refmac/file/text()")[0]

def get_seq_alias(root):
    return root.xpath("//chains/chain/seq_alias/text()")

def get_chains(root):
    chains = []
    for c in root.xpath("//chains/chain"):
        alias = c.xpath("seq_alias/text()")[0]
        ch = Chain(c.xpath("name/text()")[0], seq_alias_map[alias], alias)
        ch.read_fasta(root)
        chains.append(ch)

    return chains

def get_entities(chains):
    entities = []

    for c in chains:
        done = False

        for e in entities:
            if c.entity_id == e.entity_id:
                e.chains.append(c)
                done = True
                break

        if not done:
            ent = Entity(c.entity_id, c.seq, c)
            entities.append(ent)

    return entities

def fix_entity(mmcif, entities):
    cat = get_category(mmcif, "entity")

    if cat is None:
        return False

    idx = 0
    types = cat.getItem("type").value

    # Remove the polymers
    while idx < len(types):
        if types[idx] == "polymer":
            del cat.getItem("type").value[idx]
            del cat.getItem("id").value[idx]
        else:
            idx += 1

    for e in entities:
        cat.getItem("id").setValue(e.entity_id)
        cat.getItem("type").setValue("polymer")

    return True

def fix_entity_poly_seq(mmcif, entities):
    cat = get_category(mmcif, "entity_poly_seq")

    if cat is None:
        return False

    db = cat.parent
    cat.remove()

    cat = db.setCategory("entity_poly_seq")
    cat.setItem("entity_id")
    cat.setItem("num")
    cat.setItem("mon_id")
    cat.setItem("hetero")

    brackets = False
    buff = ""

    for e in entities:
        idx = 1
        for r in e.seq:
            if r == "(":
                brackets = True
                buff = ""
                continue
            if r == ")":
                brackets = False
                cat.getItem("mon_id").setValue(buff)
                cat.getItem("hetero").setValue("y")
                cat.getItem("entity_id").setValue(e.entity_id)
                cat.getItem("num").setValue(idx)
                idx += 1
                continue
            if brackets:
                buff += r
                continue
            if r in residue_map_1to3:
                cat.getItem("mon_id").setValue(residue_map_1to3[r])
                cat.getItem("hetero").setValue("n")
                cat.getItem("entity_id").setValue(e.entity_id)
                cat.getItem("num").setValue(idx)
            else:
                traceback.print_stack()
                print "Residue UNKNOWN!"
                sys.exit(1)

            idx += 1

    return True

def add_entity_poly(mmcif, entities):
    cat = get_category(mmcif, "entity_poly")

    # Remove if it exists
    if cat is not None:
        cat.remove()

    # Add in the same block as entity
    db = get_category(mmcif, "entity").parent
    cat = db.setCategory("entity_poly")
    cat.setItem("entity_id")
    cat.setItem("pdbx_seq_one_letter_code_can")
    cat.setItem("pdbx_seq_one_letter_code")

    for i in range(len(entities)):
        cat.getItem("entity_id").setValue(entities[i].entity_id)
        can = ""

        brackets = False

        for r in entities[i].seq:
            if r == '(':
                can += "X"
                brackets = True
                continue
            if r == ')':
                brackets = False
                continue
            if brackets:
                continue
            if r in residue_map_1to3:
                can += r
            # TODO: To be fixed once I have the residue library
            else:
                traceback.print_stack()
                print "Residue UNKNOWN!"
                sys.exit(1)

        cat.getItem("pdbx_seq_one_letter_code_can").setValue("\n".join(re.findall(".{80}|.+$", can)))
        cat.getItem("pdbx_seq_one_letter_code").setValue("\n".join(re.findall(".{80}|.+$", entities[i].seq)))

def add_pdbx_non_poly_scheme(mmcif):
    cat = get_category(mmcif, "pdbx_non_poly_scheme")

    # Remove if it exists
    if cat is not None:
        cat.remove()

    # Add in the same block as entity
    db = get_category(mmcif, "entity").parent
    cat = db.setCategory("pdbx_non_poly_scheme")
    cat.setItem("asym_id")
    cat.setItem("entity_id")
    cat.setItem("mon_id")
    cat.setItem("ndb_seq_num")
    cat.setItem("pdb_seq_num")
    cat.setItem("auth_seq_num")
    cat.setItem("pdb_mon_id")
    cat.setItem("auth_mon_id")
    cat.setItem("pdb_strand_id")
    cat.setItem("pdb_ins_code")

    _, nonpoly = split_entities(mmcif)

    atom = get_category(mmcif, "atom_site")
    ent = atom.getItem("label_entity_id")

    for idx, e in enumerate(ent.value):
        if int(e) in nonpoly:
            cat.getItem("asym_id").setValue(atom.getItem("label_asym_id").value[idx])
            cat.getItem("entity_id").setValue(atom.getItem("label_entity_id").value[idx])
            cat.getItem("mon_id").setValue(atom.getItem("label_comp_id").value[idx])
            cat.getItem("ndb_seq_num").setValue(atom.getItem("label_seq_id").value[idx])
            cat.getItem("pdb_seq_num").setValue(atom.getItem("pdbx_PDB_residue_no").value[idx])
            cat.getItem("auth_seq_num").setValue(atom.getItem("auth_seq_id").value[idx])
            cat.getItem("pdb_mon_id").setValue(atom.getItem("pdbx_PDB_residue_name").value[idx])
            cat.getItem("auth_mon_id").setValue(atom.getItem("auth_comp_id").value[idx])
            cat.getItem("pdb_strand_id").setValue(atom.getItem("auth_asym_id").value[idx])
            cat.getItem("pdb_ins_code").setValue(atom.getItem("pdbx_PDB_ins_code").value[idx])

if __name__ == '__main__':
    xml = sys.argv[1]

    if not os.path.isfile(xml):
        print "The XML file doesn't exist"
        sys.exit(1)

    outfname = os.path.basename(xml).rsplit(os.path.extsep, 1)[0] + "_out.cif"

    parser = etree.XMLParser(ns_clean=True, remove_blank_text=True)
    doc = etree.parse(xml, parser).getroot()

    dictionary = get_dictionary()

    refmac_fname = get_refmac(doc)
    refmac = get_mmcif(refmac_fname)

    # Get the entity_ids for the seq_aliases (i.e. polymers to be replaced)
    seq_alias_map = get_seq_alias_map(get_seq_alias(doc), refmac)
    chains = get_chains(doc)
    entities = get_entities(chains)

    for e in entities:
        print e

    fix_entity(refmac, entities)
    fix_entity_poly_seq(refmac, entities)
    add_entity_poly(refmac, entities)
    add_pdbx_non_poly_scheme(refmac)

    for e in entities:
        print get_atom_site_seq(refmac, e.entity_id)

    outfile = mmcifIO.CifFileWriter(outfname)
    outfile.write(refmac)