import_v1_data.py

import glob

from django.utils import timezone
from django.core.management import BaseCommand, CommandError
import mysql.connector
from pybel import readfile

from ippidb.models import Bibliography, Protein, Taxonomy, MolecularFunction, \
    Domain, ProteinDomainBoundComplex, ProteinDomainPartnerComplex, Symmetry, Ppi, PpiComplex, Disease, \
    Compound, MDDRCompoundImport, MDDRActivityClass, TestActivityDescription, CellLine


class MyConverter(mysql.connector.conversion.MySQLConverter):

    def row_to_python(self, row, fields):
        row = super(MyConverter, self).row_to_python(row, fields)

        def to_unicode(col):
            if type(col) == bytearray:
                return col.decode('utf-8')
            return col

        return[to_unicode(col) for col in row]


class Command(BaseCommand):

    help = "Import data from the local v1 database"

    def add_arguments(self, parser):
        parser.add_argument(
            '--bibliographies',
            action='store_true',
            dest='bibliographies',
            default=False,
            help='Flush and migrate bibliographies',
        )
        parser.add_argument(
            '--proteins',
            action='store_true',
            dest='proteins',
            default=False,
            help='Flush and migrate proteins',
        )
        parser.add_argument(
            '--domains',
            action='store_true',
            dest='domains',
            default=False,
            help='Flush and migrate domains',
        )
        parser.add_argument(
            '--symmetries',
            action='store_true',
            dest='symmetries',
            default=False,
            help='Flush and create symmetries',
        )
        parser.add_argument(
            '--ppi',
            action='store_true',
            dest='ppi',
            default=False,
            help='Flush and migrate ppis and complexes',
        )
        parser.add_argument(
            '--mddr',
            action='store_true',
            dest='mddr',
            default=False,
            help='Flush and import MDDR database',
        )
        parser.add_argument(
            '--compound',
            action='store_true',
            dest='compound',
            default=False,
            help='Flush and migrate compounds',
        )
        parser.add_argument(
            '--test-activity-description',
            action='store_true',
            dest='test-activity-description',
            default=False,
            help='Flush and migrate test activity descriptions',
        )
        parser.add_argument(
            '--stoponfail',
            action='store_true',
            dest='stoponfail',
            default=False,
            help='Stop on fail',
        )

    def handle(self, *args, **options):
        conn = mysql.connector.connect(
            converter_class=MyConverter, host="localhost", user="root", password="ippidb", database="ippidb")
        cursor = conn.cursor()
        if options['bibliographies']:
            cursor.execute("""SELECT * FROM biblio""")
            rows = cursor.fetchall()
            Bibliography.objects.all().delete()
            self.stdout.write(
                self.style.SUCCESS('Successfully flushed bibliography table'))
            for row in rows:
                try:
                    b = Bibliography()
                    if row[1] == 'article':
                        b.source = 'PM'
                    else:
                        b.source = 'PT'
                    b.id_source = row[2]
                    b.save(autofill=True)
                except Exception as e:
                    self.stdout.write(
                        self.style.ERROR('Failed inserting {}'.format(row[2])))
                else:
                    self.stdout.write(
                        self.style.SUCCESS('Successfully inserted {}'.format(row[2])))
        if options['proteins']:
            cursor.execute("""SELECT * FROM protein""")
            rows = cursor.fetchall()
            Protein.objects.all().delete()
            self.stdout.write(
                self.style.SUCCESS('Successfully flushed protein table'))
            Taxonomy.objects.all().delete()
            self.stdout.write(
                self.style.SUCCESS('Successfully flushed taxonomy table'))
            MolecularFunction.objects.all().delete()
            self.stdout.write(
                self.style.SUCCESS('Successfully flushed molecular function table'))
            for row in rows:
                try:
                    p = Protein()
                    p.uniprot_id = row[1]
                    p.save(autofill=True)
                except Exception as e:
                    if options['stoponfail']:
                        import traceback
                        self.stderr.write(traceback.format_exc())
                        raise CommandError(
                            'Failed inserting {} {}'.format(row[1], row[2]))
                    else:
                        self.stdout.write(
                            self.style.ERROR('Failed inserting {} {}'.format(row[1], row[2])))
                else:
                    self.stdout.write(
                        self.style.SUCCESS('Successfully inserted {} {}'.format(row[1], row[2])))
        if options['domains']:
            cursor.execute("""SELECT * FROM domain""")
            rows = cursor.fetchall()
            Domain.objects.all().delete()
            self.stdout.write(
                self.style.SUCCESS('Successfully flushed domain table'))
            for row in rows:
                try:
                    p = Domain()
                    p.pfam_acc = row[2]
                    p.domain_family = row[4]
                    p.save(autofill=True)
                except Exception as e:
                    if options['stoponfail']:
                        import traceback
                        self.stderr.write(traceback.format_exc())
                        raise CommandError(
                            'Failed inserting {} {}'.format(row[1], row[2]))
                    else:
                        self.stdout.write(
                            self.style.ERROR('Failed inserting {} {}'.format(row[1], row[2])))
                else:
                    self.stdout.write(
                        self.style.SUCCESS('Successfully inserted {} {}'.format(row[1], row[2])))
        if options['symmetries']:
            Symmetry.objects.all().delete()
            self.stdout.write(
                self.style.SUCCESS('Successfully flushed symmetries table'))
            rows = [
                ['AS', 'asymmetric'],
                    ['C2', 'C2 symmetry'],
                    ['D2', 'D2 symmetry'],
                    ['C3', 'C3 symmetry'],
                    ['D3', 'D3 symmetry'],
                    ['C4', 'C4 symmetry'],
                    ['D4', 'D4 symmetry'],
                    ['C5', 'C5 symmetry'],
                    ['D5', 'D5 symmetry'],
            ]
            for row in rows:
                try:
                    symmetry = Symmetry()
                    symmetry.code = row[0]
                    symmetry.description = row[1]
                    symmetry.save()
                except Exception as e:
                    if options['stoponfail']:
                        import traceback
                        self.stderr.write(traceback.format_exc())
                        raise CommandError(
                            'Failed inserting {} {}'.format(row[0], row[1]))
                    else:
                        self.stdout.write(
                            self.style.ERROR('Failed inserting {} {}'.format(row[0], row[1])))
                else:
                    self.stdout.write(
                        self.style.SUCCESS('Successfully inserted {} {}'.format(row[0], row[1])))
        if options['ppi']:
            sql_request_string = '''
select distinct protein.NumUniprot, domain.PfamNumAccession, complexe.NbCopy, cmpdAction.IDComplexeBound, bindingSiteEvidence.CodePDB, 'part1', ppi.IDPPI, disease.Disease  from bindingSite inner join ppi on (bindingSite.IDBindingSite=ppi.IDBindingSite1) inner join complexe on (ppi.IDComplexe1=complexe.IDComplexe) left outer join cmpdAction on (complexe.IDComplexe=cmpdAction.IDComplexeBound) inner join protein on (bindingSite.IDProtein=protein.IDProtein) inner join domain on (bindingSite.IDDomain=domain.IDDomain) inner join disease on (disease.IDPPI=ppi.IDPPI) left outer join bindingSiteEvidence on (ppi.IDPPI=bindingSiteEvidence.IDPPI)
union
select distinct protein.NumUniprot, domain.PfamNumAccession  , complexe.NbCopy, cmpdAction.IDComplexeBound, null, 'part2', ppi.IDPPI, disease.Disease  from bindingSite inner join ppi on (bindingSite.IDBindingSite=ppi.IDBindingSite2) inner join complexe on (ppi.IDComplexe2=complexe.IDComplexe) left outer join cmpdAction on (complexe.IDComplexe=cmpdAction.IDComplexeBound) inner join protein on (bindingSite.IDProtein=protein.IDProtein) inner join domain on (bindingSite.IDDomain=domain.IDDomain) inner join disease on (disease.IDPPI=ppi.IDPPI)'''
            cursor.execute(sql_request_string)
            rows = cursor.fetchall()
            ProteinDomainBoundComplex.objects.all().delete()
            ProteinDomainPartnerComplex.objects.all().delete()
            Disease.objects.all().delete()
            Ppi.objects.all().delete()
            PpiComplex.objects.all().delete()
            self.stdout.write(
                self.style.SUCCESS('Successfully flushed protein domain complex, PPI, and disease tables'))
            ppi_ids_mapping = {}
            for row in rows:
                try:
                    # create or retrieve Ppi object
                    if row[5] == 'part1':
                        ppi = Ppi()
                        disease, created = Disease.objects.get_or_create(
                            name=row[7])
                        ppi.pdb_id = row[4]
                        ppi.pockets_nb = 1
                        ppi.symmetry = Symmetry.objects.get(code='AS')
                        ppi.save()
                        ppi.diseases.add(disease)
                        ppi.save()
                        ppi_ids_mapping[row[6]] = ppi.id
                    else:
                        ppi = Ppi.objects.get(id=ppi_ids_mapping[row[6]])
                    # create a complex
                    if row[3] is None:
                        c = ProteinDomainPartnerComplex()
                    else:
                        c = ProteinDomainBoundComplex()
                    protein = Protein.objects.get(uniprot_id=row[0])
                    c.protein = protein
                    domain = Domain.objects.get(pfam_acc=row[1])
                    c.domain = domain
                    c.ppc_copy_nb = row[2]
                    if isinstance(c, ProteinDomainBoundComplex):
                        c.ppp_copy_nb_per_p = 1
                    c.save()
                    # create the PpiComplex object
                    ppi_complex = PpiComplex()
                    ppi_complex.ppi = ppi
                    ppi_complex.complex = c
                    ppi_complex.cc_nb = 1
                    ppi_complex.save()
                except Exception as e:
                    if options['stoponfail']:
                        import traceback
                        self.stderr.write(traceback.format_exc())
                        raise CommandError(
                            'Failed inserting {} {}'.format(row[0], row[1]))
                    else:
                        self.stdout.write(
                            self.style.ERROR('Failed inserting {} {}'.format(row[0], row[1])))
                else:
                    self.stdout.write(
                        self.style.SUCCESS('Successfully inserted {} {}'.format(row[0], row[1])))
            cursor.execute("""SELECT * FROM testActivityDescription""")
            rows = cursor.fetchall()
            TestActivityDescription.objects.all().delete()
            CellLine.objects.all().delete()
            self.stdout.write(
                self.style.SUCCESS('Successfully flushed test activity descriptions table and cell lines table'))
            for row in rows:
                try:
                    tad = TestActivityDescription()
                    cursor.execute(
                        """select IDSource from biblio where IDBiblio={}""".format(row[2]))
                    biblio_row = cursor.fetchone()
                    biblio = Bibliography.objects.get(id_source=biblio_row[0])
                    tad.biblio = biblio
                    tad.ppi = Ppi.objects.get(id=ppi_ids_mapping[row[3]])
                    tad.test_name = row[4]
                    tad.test_type = row[7].upper()
                    tad.test_modulation_type = row[8][0]
                    tad.nb_active_compounds = row[9]
                    if row[16] is not None:
                        tad.cell_line, created = CellLine.objects.get_or_create(
                            name=row[16])
                    tad.save()
                except Exception as e:
                    if options['stoponfail']:
                        import traceback
                        self.stderr.write(traceback.format_exc())
                        raise CommandError(
                            'Failed inserting {} {}'.format(row[1], row[2]))
                    else:
                        self.stdout.write(
                            self.style.ERROR('Failed inserting {} {}'.format(row[1], row[2])))
                else:
                    self.stdout.write(
                        self.style.SUCCESS('Successfully inserted {}'.format(row[2])))
        if options['mddr']:
            MDDRCompoundImport.objects.all().delete()
            MDDRActivityClass.objects.all().delete()
            self.stdout.write(
                self.style.SUCCESS('Successfully flushed MDDR Compound and Activity class tables'))
            for sdf_file in glob.glob('/home/hmenager/iPPIDB/mddr20151_2d.sdf/*.sdf'):
                for item in readfile("sdf", sdf_file):
                    try:
                        m = MDDRCompoundImport()
                        m.mddr_name = item.data['MOLNAME']
                        m.canonical_smile = str(item)
                        m.dvpmt_phase = item.data['PHASE']
                        m.db_import_date = timezone.now()
                        m.save()
                        for activity_class_name in item.data['ACTIV_CLASS'].split(','):
                            activity_class, created = MDDRActivityClass.objects.get_or_create(
                                name=activity_class_name)
                            m.activity_classes.add(activity_class)
                        m.save()
                    except Exception as e:
                        if options['stoponfail']:
                            import traceback
                            self.stderr.write(traceback.format_exc())
                            raise CommandError(
                                'Failed inserting {}'.format(str(item)))
                        else:
                            self.stdout.write(
                                self.style.ERROR('Failed inserting {}'.format(str(item))))
                    else:
                        self.stdout.write(
                            self.style.SUCCESS('Successfully inserted {} {}'.format(item.data['MOLNAME'], str(item))))
        if options['compound']:
            cursor.execute("""SELECT * FROM compound""")
            rows = cursor.fetchall()
            Compound.objects.all().delete()
            self.stdout.write(
                self.style.SUCCESS('Successfully flushed compound table'))
            for row in rows:
                try:
                    compound = Compound()
                    compound.canonical_smile = row[1]
                    compound.is_macrocycle = (row[4] == 'Y')
                    compound.aromatic_ratio = row[5]
                    compound.balaban_index = row[6]
                    compound.fsp3 = row[7]  # Csp3Ratio
                    compound.gc_molar_refractivity = row[
                        10]  # GCMolarRefractivity
                    compound.log_d = row[13]  # LogD
                    compound.a_log_p = row[14]  # ALogP
                    compound.mean_atom_vol_vdw = row[15]  # MeanAtomVolVdW
                    compound.molecular_weight = row[16]  # MolecularWeight
                    compound.nb_acceptor_h = row[17]  # NbAcceptorH
                    compound.nb_aliphatic_amines = row[
                        18]  # NbAliphaticsAmines
                    compound.nb_aromatic_bonds = row[19]  # NbAromaticBonds
                    compound.nb_aromatic_ether = row[20]  # NbAromaticsEther
                    compound.nb_aromatic_sssr = row[21]  # NbAromaticsSSSR
                    compound.nb_atom = row[22]  # NbAtom
                    compound.nb_atom_non_h = row[23]  # NbAtomNonH
                    compound.nb_benzene_like_rings = row[24]  # NbBenzLikeRings
                    compound.nb_bonds = row[25]  # NbBonds
                    compound.nb_bonds_non_h = row[26]  # NbBondsNonH
                    compound.nb_br = row[27]  # NbBr
                    compound.nb_c = row[28]  # NbC
                    compound.nb_chiral_centers = row[29]  # NbChiralCenters
                    compound.nb_circuits = row[30]  # NbCircuits
                    compound.nb_cl = row[31]  # NbCl
                    compound.nb_csp2 = row[32]  # NbCsp2
                    compound.nb_csp3 = row[33]  # NbCsp3
                    compound.nb_donor_h = row[34]  # NbDonorH
                    compound.nb_double_bonds = row[35]  # NbDoubleBonds
                    compound.nb_f = row[36]  # NbF
                    compound.nb_i = row[37]  # NbI
                    compound.nb_multiple_bonds = row[38]  # NbMultBonds
                    compound.nb_n = row[39]  # NbN
                    compound.nb_o = row[40]  # NbO
                    compound.nb_rings = row[41]  # NbRings
                    compound.nb_rotatable_bonds = row[42]  # NbRotatableBonds
                    compound.randic_index = row[44]  # RandicIndex
                    compound.rdf070m = row[45]  # RDF070m
                    compound.rotatable_bond_fraction = row[
                        46]  # RotatableBondFraction
                    compound.sum_atom_polar = row[47]  # SumAtomPolar
                    compound.sum_atom_vol_vdw = row[48]  # SumAtomVolVdW
                    compound.tpsa = row[51]  # TPSA
                    compound.ui = row[52]  # Ui
                    compound.wiener_index = row[54]  # WienerIndex
                    if row[55] != 'N':
                        compound.common_name = row[55]  # CmpdNameSh
                    compound.pubchem_id = row[56]  # IdPubchem
                    if row[57] != 'N':
                        compound.chemspider_id = row[57]  # IdPubchem
                    compound.chembl_id = row[58]
                    compound.iupac_name = row[59]
                    # compound.mddr_compound = row[60]
                    # FIXME lookup MDDR to get the right ID
                    compound.save()
                except Exception as e:
                    if options['stoponfail']:
                        import traceback
                        self.stderr.write(traceback.format_exc())
                        raise CommandError(
                            'Failed inserting {} {}'.format(row[1], row[2]))
                    else:
                        self.stdout.write(
                            self.style.ERROR('Failed inserting {} {}'.format(row[1], row[2])))
                else:
                    self.stdout.write(
                        self.style.SUCCESS('Successfully inserted {} {}'.format(row[1], row[2])))