import hashlib
import numpy as np
import os
import re
import requests
import tarfile
import time
import unittest

from absl import logging
from typing import List, NoReturn, Tuple

class MMSeqs2Runner:

  r"""Runner object

  Fetches sequence alignment and templates from MMSeqs2 server
  Based on the function run_mmseqs2 from ColabFold (sokrypton/ColabFold)
  Version 62d7558c91a9809712b022faf9d91d8b183c328c

  Relevant publications
  ----------
  * "Clustering huge protein sequence sets in linear time"
    https://doi.org/10.1038/s41467-018-04964-5
  * "MMseqs2 enables sensitive protein sequence searching for the analysis
    of massive data sets"
    https://doi.org/10.1038/nbt.3988

  Private variables
  ----------
  self.job: Job ID (five-char string)
  self.seq: Sequence to search
  self.host_url: URL address to ping for data
  self.t_url: URL address to ping for templates from PDB
  self.n_templates = Number of templates to fetch (default=20)
  self.path: Path to use
  self.tarfile: Compressed file archive to download
  """

  def __init__(
      self,
      job: str,
      seq: str,
      host_url: str = "https://a3m.mmseqs.com",
      t_url: str = "https://a3m-templates.mmseqs.com/template",
      path_suffix: str = "env",
      n_templates: int = 20
    ):

    r"""Initialize runner object

    Parameters
    ----------
    job : Job name
    seq : Amino acid sequence
    host_url : Website to ping for sequence data
    t_url : Website to ping for template info
    path_suffix : Suffix for path info

    """

    # Clean up sequence
    self.seq = self._cleanseq( seq.upper() )

    # Come up with unique job ID for MMSeqs
    self.job = self._define_jobname( job )

    # Save everything else
    self.host_url = host_url
    self.t_url = t_url
    self.n_templates = n_templates

    self.path = "_".join( ( self.job, path_suffix ) )

    if not os.path.isdir( self.path ):
      os.system( f"mkdir { self.path }" )

    self.tarfile = f'{ self.path }/out.tar.gz'

  def _cleanseq(
      self,
      seq
    ) -> str:

    r""" Cleans the sequence to remove whitespace and noncanonical letters

    Parameters
    ----------
    seq : Amino acid sequence (only all 20 here)

    Returns
    ----------
    Cleaned up amin acid sequence

    """

    if any( [ aa in seq for aa in "BJOUXZ" ] ):
      logging.warning( "Sequence contains non-canonical amino acids!" )
      logging.warning( "Removing B, J, O, U, X, and Z from sequence" )
      seq = re.sub( r'[BJOUXZ]', '', seq )

    return re.sub( r'[^A-Z]', '', "".join( seq.split() ) )

  def _define_jobname(
      self,
      job: str
    ) -> str:

    r""" Provides a unique five-digit identifier for the job name

    Parameters
    ----------
    job : Job name

    Returns
    ----------
    Defined job name

    """

    return "_".join( (
        re.sub( r"\W+", "", "".join( job.split() ) ),
        hashlib.sha1( self.seq.encode() ).hexdigest()[ :5 ] )
      )

  def _submit(
      self
    ) -> dict:

    r"""Submit job to MMSeqs2 server

    Parameters
    ----------
    None

    Returns
    ----------
    None

    """

    data = { 'q': f">101\n{ self.seq }", 'mode': "env" }

    res = requests.post( f'{ self.host_url }/ticket/msa', data=data )
    
    try:
      out = res.json()

    except ValueError:
      out = { "status": "UNKNOWN" }

    return out

  def _status(
      self,
      idx: str
    ) -> dict:

    r"""Check status of job 

    Parameters
    ----------
    idx : Index assigned by MMSeqs2 server

    Returns
    ----------
    None

    """

    res = requests.get( f'{ self.host_url }/ticket/{ idx }' )
    
    try:
      out = res.json()
    
    except ValueError:
      out = { "status": "UNKNOWN" }
    
    return out

  def _download(
      self,
      idx: str,
      path: str
    ) -> NoReturn:

    r"""Download job outputs

    Parameters
    ----------
    idx : Index assigned by MMSeqs2 server
    path : Path to download data

    Returns
    ----------
    None
    
    """

    res = requests.get( f'{ self.host_url }/result/download/{ idx }' )
    
    with open( path, "wb" ) as out:
      out.write( res.content )

  def _search_mmseqs2(
      self
    ) -> NoReturn:

    r"""Run the search and download results
    Heavily modified from ColabFold

    Parameters
    ----------
    None

    Returns
    ----------
    None
    
    """

    if os.path.isfile( self.tarfile ):
      return

    out = self._submit()

    time.sleep( 5 + np.random.randint( 0, 5 ) )
    while out[ "status" ] in [ "UNKNOWN", "RATELIMIT" ]:
      # resubmit
      time.sleep( 5 + np.random.randint( 0, 5 ) )
      out = self._submit()

    logging.debug( f"ID: { out[ 'id' ] }" )

    while out[ "status" ] in [ "UNKNOWN", "RUNNING", "PENDING" ]:
      time.sleep( 5 + np.random.randint( 0, 5 ) )
      out = self._status( out[ "id" ] )

    if out[ "status" ] == "COMPLETE":
      self._download( out[ "id" ], self.tarfile )

    elif out[ "status" ] == "ERROR":
      raise RuntimeError( " ".join( ( "MMseqs2 API is giving errors.",
        "Please confirm your input is a valid protein sequence.",
        "If error persists, please try again in an hour." ) ) )

  def process_templates(
      self,
      templates: List[ str ] = []
    ) -> str:
    
    r"""Process templates and fetch from MMSeqs2 server

    Parameters
    ----------
    use_templates : True/False whether to use templates
    max_templates : Maximum number of templates to use

    Returns
    ----------
    Directory containing templates (empty if not using templates)

    """

    path = f"{ self.job }_env/templates_101"
    if os.path.isdir( path ):
      os.system( f"rm { path }" )

    #templates = {}
    logging.info( "\t".join( ( "seq", "pdb", "cid", "evalue" ) ) )

    pdbs = []
    with open( f"{ self.path }/pdb70.m8", "r" ) as infile:

      for line in infile:

        sl = line.rstrip().split()
        pdb = sl[ 1 ]
        if pdb in templates:
          pdbs.append( sl[ 1 ] )
          logging.info( f"{ sl[ 0 ] }\t{ sl[ 1 ] }\t{ sl[ 2 ] }\t{ sl[ 10 ] }" )
      
    if len( pdbs ) == 0:
      logging.warning( "No templates found." )
      return ""
    
    else:

      if not os.path.isdir( path ):
        os.mkdir( path )
      
      pdbs = [ t for t in pdbs if t in templates ]

      if len( templates ) == 0 or len( pdbs ) == 0:
        pdbs = ",".join( templates[ :self.n_templates ] )
      else:
        pdbs = ",".join( pdbs[ :self.n_templates ] )
      
      os.system( f"curl -v { self.t_url }/{ pdbs }" f" | tar xzf - -C { path }/" )
      
      os.system( f"cp { path }/pdb70_a3m.ffindex { path }/pdb70_cs219.ffindex" )
      
      os.system( f"touch { path }/pdb70_cs219.ffdata" )

      return path

  def _process_alignment(
      self,
      a3m_files: list,
      templates: List[ str ] = []
    ) -> Tuple[ str, str ]:
    
    r""" Process sequence alignment
    (modified from ColabFold)

    Parameters
    ----------
    a3m_files : List of files to parse
    token : Token to look for when parsing

    Returns
    ----------
    Tuple with [0] string with alignment, and [1] path to template

    """

    a3m_lines = ""

    for a3m_file in a3m_files:
      for line in open( os.path.join( self.path, a3m_file ), "r" ):
        if len( line ) > 0:
          a3m_lines += line.replace( "\x00", "" )

    return a3m_lines, self.process_templates( templates )

  def run_job(
      self,
      templates: List[ str ] = []
    ) -> Tuple[ str, str ]:
    
    r"""
    Run sequence alignments using MMseqs2

    Parameters
    ----------
    use_templates: Whether to use templates

    Returns
    ----------
    Tuple with [0] string with alignment, and [1] path to template

    """

    self._search_mmseqs2()

    a3m_files = [ "uniref.a3m", "bfd.mgnify30.metaeuk30.smag30.a3m" ]

    # extract a3m files
    if not os.path.isfile( os.path.join( self.path, a3m_files[ 0 ] ) ):
      with tarfile.open( self.tarfile ) as tar_gz:
        tar_gz.extractall( self.path )  

    return self._process_alignment( a3m_files, templates )

###########################
# TESTS

class TestMMseqs2Runner( unittest.TestCase ):

  r""" Test object that verifies that everything works as expected

  """

  def test_seq_checker( self ):

    seq = ( "MNIFEMLRIDEGLRLKIYKDTEGYYTIGIGHLLTKSPSLNAAKSELDKAIGRNCNGVIT"
            "KDEAEKLFNQDVDAAVRGILRNAKLKPVYDSLDAVRRCALINMVFQMGETGVAGFTNSL"
            "RMLQQKRWDEAAVNLAKSRWYNQTPNRAKRVITTFRTGTWDAYKNL" )

    job = "T4_lysozyme"

    runner = MMSeqs2Runner( job, seq )
    assert runner.seq == seq

    runner_2 = MMSeqs2Runner( job, seq + " BJO\tUXZ" )
    assert runner_2.seq == seq

  def test_jobname( self ):

    seq = ( "MNIFEMLRIDEGLRLKIYKDTEGYYTIGIGHLLTKSPSLNAAKSELDKAIGRNCNGVIT"
            "KDEAEKLFNQDVDAAVRGILRNAKLKPVYDSLDAVRRCALINMVFQMGETGVAGFTNSL"
            "RMLQQKRWDEAAVNLAKSRWYNQTPNRAKRVITTFRTGTWDAYKNL" )

    job = "T4_lysozyme"

    runner = MMSeqs2Runner( job, seq )
    assert runner.job == f"{ job }_bdd05"

  def test_db( self ):

    # Here we fetch the actual data
    # Need to somehow do this without making a new directory
    pass

  def test_msa( self ):

    # Here the ability to convert the MSA file to a sinlg string is tested
    pass

  def test_templates( self ):

    # Here the ability to fetch multiple templates is tested
    # Also test that certain templates are not fetched
    pass

if __name__ == "__main__":
  logging.set_verbosity( logging.INFO )
  unittest.main()