package jalview.ws2.actions.secstructpred;

import java.io.IOException;
import java.util.List;

import jalview.analysis.AlignSeq;
import jalview.analysis.AlignmentAnnotationUtils;
import jalview.analysis.SeqsetUtils;
import jalview.analysis.SeqsetUtils.SequenceInfo;
import jalview.api.AlignViewportI;
import jalview.bin.Console;
import jalview.commands.RemoveGapsCommand;
import jalview.datamodel.Alignment;
import jalview.datamodel.AlignmentAnnotation;
import jalview.datamodel.AlignmentI;
import jalview.datamodel.AlignmentView;
import jalview.datamodel.SeqCigar;
import jalview.datamodel.SequenceI;
import jalview.io.JPredFile;
import jalview.io.JnetAnnotationMaker;
import jalview.util.Comparison;
import jalview.util.MessageManager;
import jalview.ws.params.ArgumentI;
import jalview.ws2.actions.BaseJob;
import jalview.ws2.actions.BaseTask;
import jalview.ws2.actions.ServiceInputInvalidException;
import jalview.ws2.api.Credentials;
import jalview.ws2.api.JobStatus;
import jalview.ws2.client.api.SecStructPredWebServiceClientI;

public class SecStructPredPDBSearchTask extends
        BaseTask<SecStructPredPDBSearchTask.SecStructPredJob, AlignmentI>
{
  private final SecStructPredWebServiceClientI client;

  private final AlignmentView alignmentView;
  private final AlignmentI currentView;
  private final char gapChar;

  SecStructPredPDBSearchTask(SecStructPredWebServiceClientI client,
          List<ArgumentI> args, Credentials credentials,
          AlignViewportI viewport)
  {
    super(client, args, credentials);
    this.client = client;
    this.alignmentView = viewport.getAlignmentView(true);
    this.currentView = viewport.getAlignment();
    this.gapChar = viewport.getGapCharacter();
  }

  @Override
  protected List<SecStructPredJob> prepareJobs()
          throws ServiceInputInvalidException
  {
    SeqCigar[] msf = alignmentView.getSequences();
    if (msf.length > 1)
      throw new ServiceInputInvalidException(MessageManager.getString(
              "error.implementation_error_multiple_single_sequence_prediction_jobs_not_supported"));
    SequenceI seq = msf[0].getSeq('-');
    if (seq == null)
      throw new ServiceInputInvalidException("Missing sequence.");
    int[] delMap = alignmentView.getVisibleContigMapFor(seq.gapMap());
    var seqInfo = SeqsetUtils.SeqCharacterHash(seq);
    seq.setSequence(AlignSeq.extractGaps(Comparison.GapChars,
            alignmentView.getASequenceString('-', 0)));
    if (seq.getLength() < 20)
      throw new ServiceInputInvalidException(
              "Sequence is too short to predict with JPred - need at least 20 amino acids.");
    var job = new SecStructPredJob(seq, delMap, seqInfo);
    job.setStatus(JobStatus.READY);
    return List.of(job);
  }
  
  public static final int MSA_INDEX = 0;

  @Override
  protected AlignmentI collectResult(List<SecStructPredJob> jobs)
          throws IOException
  {
    var job = jobs.get(0);  // There should be exactly one job
    var status = job.getStatus();
    Console.info(
            String.format("sec str pred job \"%s\" finished with status %s",
                    job.getServerJob().getJobId(), status));
    JPredFile predictionFile = client.getPredictionFile(job.getServerJob());
    SequenceI[] preds = predictionFile.getSeqsAsArray();
    Alignment aln = new Alignment(preds);
    int queryPosition = predictionFile.getQuerySeqPosition();
    SequenceI profileSeq = aln.getSequenceAt(queryPosition);
    if (job.delMap != null)
    {
      SequenceI[] seqs = (SequenceI[]) alignmentView.getAlignmentAndHiddenColumns(gapChar)[0];
      if (MSA_INDEX >= seqs.length)
        throw new Error(MessageManager.getString(
                "error.implementation_error_invalid_msa_index_for_job"));
      new RemoveGapsCommand(
              MessageManager.getString("label.remove_gaps"),
              new SequenceI[] { seqs[MSA_INDEX] }, currentView);
      profileSeq.setSequence(seqs[MSA_INDEX].getSequenceAsString());
    }
    if (!SeqsetUtils.SeqCharacterUnhash(aln.getSequenceAt(queryPosition), job.info))
      throw new IOException(MessageManager.getString("exception.couldnt_recover_sequence_props_for_jnet_query"));
    aln.setDataset(currentView.getDataset());
    try
    {
      JnetAnnotationMaker.add_annotation(predictionFile, aln, queryPosition, true,
              job.delMap);
    } catch (Exception e)
    {
      throw new IOException(e);
    }
    alignToProfileSequence(aln, profileSeq);
    if (job.delMap != null)
      aln.setHiddenColumns(aln.propagateInsertions(profileSeq, alignmentView));

    for (AlignmentAnnotation alnAnnot : aln.getAlignmentAnnotation())
    {
      if (alnAnnot.sequenceRef != null)
      {
        AlignmentAnnotationUtils.replaceAnnotationOnAlignmentWith(alnAnnot,
                alnAnnot.label, getClass().getSimpleName());
      }
    }
    aln.setSeqrep(profileSeq);
    return aln;
  }
  
  /**
   * Given an alignment where all other sequences except profileseq are
   * aligned to the ungapped profileseq, insert gaps in the other sequences to
   * realign them with the residues in profileseq.
   * 
   * Shamelessly copied from JPredThread.
   * 
   * @param al
   * @param profileseq
   */
  private static void alignToProfileSequence(AlignmentI al, SequenceI profileseq)
  {
    char gc = al.getGapCharacter();
    int[] gapMap = profileseq.gapMap();
    // insert gaps into profile
    for (int lp = 0, r = 0; r < gapMap.length; r++)
    {
      if (gapMap[r] - lp > 1)
      {
        StringBuffer sb = new StringBuffer();
        for (int s = 0, ns = gapMap[r] - lp; s < ns; s++)
        {
          sb.append(gc);
        }
        for (int s = 1, ns = al.getHeight(); s < ns; s++)
        {
          String sq = al.getSequenceAt(s).getSequenceAsString();
          int diff = gapMap[r] - sq.length();
          if (diff > 0)
          {
            // pad gaps
            sq = sq + sb;
            while ((diff = gapMap[r] - sq.length()) > 0)
            {
              sq = sq + ((diff >= sb.length()) ? sb.toString()
                      : sb.substring(0, diff));
            }
            al.getSequenceAt(s).setSequence(sq);
          }
          else
          {
            al.getSequenceAt(s).setSequence(sq.substring(0, gapMap[r])
                    + sb.toString() + sq.substring(gapMap[r]));
          }
        }
      }
      lp = gapMap[r];
    }
  }

  public static class SecStructPredJob extends BaseJob
  {
    private final SequenceInfo info;

    private final int[] delMap;

    SecStructPredJob(SequenceI querySequence, int[] delMap,
            SequenceInfo info)
    {
      super(List.of(querySequence));
      this.delMap = delMap;
      this.info = info;
    }

    @Override
    public boolean isInputValid()
    {
      return true;
    }
  }

}