// Copyright 2019 Global Phasing Ltd.
//
// This program compares sequence from SEQRES and from the model.

#include <gemmi/model.hpp>
#include <gemmi/polyheur.hpp>  // for setup_entities, align_sequence_to_polymer
#include <gemmi/align.hpp>     // for align_sequence_to_polymer
#include <gemmi/seqalign.hpp>  // for align_string_sequences
#include <gemmi/read_coor.hpp>

#include <cstdio>   // for printf, fprintf, putchar
#include <cstdlib>  // for atoi
#define GEMMI_PROG align
#include "options.h"

namespace {

using std::printf;

enum OptionIndex { Match=4, Mismatch, GapOpen, GapExt,
                   CheckMmcif, PrintOneLetter, Query, Target, TextAlign, Rmsd };

const option::Descriptor Usage[] = {
  { NoOp, 0, "", "", Arg::None,
    "Pairwise sequence alignment with scoring matrix and affine gap penalty."
    "\n\nUsage:"
    "\n\n" EXE_NAME " [options] FILE[...]"
    "\n    Aligns sequence from the model to the full sequence (SEQRES)."
    "\n    Both are from the same FILE - either in the PDB or mmCIF format."
    "\n    If the mmCIF format is used, option --check-mmcif can be used."
    "\n\n" EXE_NAME " [options] --query=CHAIN1 --target=CHAIN2 FILE1 FILE2"
    "\n    Aligns CHAIN1 from FILE1 to CHAIN2 from FILE2."
    "\n    By default, the sequence of residues in the model is used."
    "\n    To use SEQRES prepend '+' to the chain name (e.g. --query=+A)."
    "\n\n" EXE_NAME " [options] --text-align STRING1 STRING2"
    "\n    Aligns two ASCII strings (used for testing)."
    "\n\nOptions:" },
  CommonUsage[Help],
  CommonUsage[Version],

  { CheckMmcif, 0, "", "check-mmcif", Arg::None,
    "  --check-mmcif  \tchecks alignment against _atom_site.label_seq_id" },
  { Query, 0, "", "query", Arg::Required,
    "  --query=[+]CHAIN  \tAlign CHAIN from file INPUT1." },
  { Target, 0, "", "target", Arg::Required,
    "  --target=[+]CHAIN  \tAlign CHAIN from file INPUT2." },
  { TextAlign, 0, "", "text-align", Arg::None,
    "  --text-align  \tAlign characters in two strings (for testing)." },

  { NoOp, 0, "", "", Arg::None, "\nScoring (absolute values):" },
  { Match, 0, "", "match", Arg::Int,
    "  --match=INT  \tMatch score (default: 1)." },
  { Mismatch, 0, "", "mism", Arg::Int,
    "  --mism=INT  \tMismatch penalty (default: -1)." },
  { GapOpen, 0, "", "gapo", Arg::Int,
    "  --gapo=INT  \tGap opening penalty (default: -1)." },
  { GapExt, 0, "", "gape", Arg::Int,
    "  --gape=INT  \tGap extension penalty (default: -1)." },

  { NoOp, 0, "", "", Arg::None, "\nOutput options:" },
  { PrintOneLetter, 0, "p", "", Arg::None,
    "  -p  \tPrint formatted alignment with one-letter codes." },
  { Rmsd, 0, "", "rmsd", Arg::None,
    "  --rmsd  \tIn addition to aligning two CHAINs (--query and --target), "
                 "superpose structures and print RMSD." },
  CommonUsage[Verbose],
  { 0, 0, 0, 0, 0, 0 }
};

void print_alignment_details(const gemmi::AlignmentResult& result,
                             const std::string& chain_name,
                             const gemmi::ConstResidueSpan& polymer,
                             const gemmi::Entity& ent) {
  std::vector<bool> gaps = prepare_free_gapo(polymer, ent.polymer_type);
  auto gap = gaps.begin();
  int seq_pos = 0;
  auto model_residues = polymer.first_conformer();
  auto res = model_residues.begin();
  for (gemmi::AlignmentResult::Item item : result.cigar) {
    char op = item.op();
    for (uint32_t i = 0; i < item.len(); ++i) {
      std::string fmon = gemmi::Entity::first_mon(ent.full_sequence[seq_pos]);
      printf("  %s ", chain_name.c_str());
      if (op == 'I' || op == 'M') {
        seq_pos++;
        printf("%4d %-3s -", seq_pos, fmon.c_str());
      } else {
        printf("         -");
      }
      if (op == 'D' || op == 'M') {
        printf(" %-3s %4d%c",
               res->name.c_str(), *res->seqid.num, res->seqid.icode);
        if (res->label_seq)
          printf("   id:%4d %c",
                 *res->label_seq, *res->label_seq == seq_pos ? ' ' : '!');
        std::putchar(*gap++ ? '^' : ' ');
        if (op == 'D' || fmon != res->name)
          printf("    <-- BAD");
        res++;
      }
      printf("\n");
    }
  }
}

void check_label_seq_id(const gemmi::AlignmentResult& result,
                        const gemmi::ConstResidueSpan& polymer) {
  int seq_pos = 1;
  auto residues = polymer.first_conformer();
  auto res = residues.begin();
  for (gemmi::AlignmentResult::Item item : result.cigar) {
    char op = item.op();
    for (uint32_t i = 0; i < item.len(); ++i) {
      if (op == 'D' || op == 'M') {
        if (*res->label_seq != seq_pos)
          printf("NOTE: %s with label_seq_id %d , expected %d.\n",
                 res->name.c_str(), *res->label_seq, seq_pos);
        res++;
      }
      if (op == 'I' || op == 'M')
        seq_pos++;
    }
  }
}

const gemmi::Model& get_first_model(gemmi::Structure& st) {
  gemmi::setup_entities(st);
  if (st.models.empty())
    gemmi::fail("No atoms found. Wrong input file?");
  return st.models[0];
}

gemmi::ConstResidueSpan get_polymer(const gemmi::Model& model,
                                    const std::string& chain_name) {
  const gemmi::Chain* ch = model.find_chain(chain_name);
  if (!ch)
    gemmi::fail("No such chain: " + chain_name);
  gemmi::ConstResidueSpan polymer = ch->get_polymer();
  if (!polymer)
    gemmi::fail("Polymer not found in chain " + chain_name);
  return polymer;
}

const gemmi::Entity* get_entity(gemmi::Structure& st,
                                const std::string& chain_name) {
  auto polymer = get_polymer(get_first_model(st), chain_name);
  if (const gemmi::Entity* ent = st.get_entity_of(polymer))
    return ent;
  gemmi::fail("No sequence (SEQRES) for chain " + chain_name);
}

void print_one_letter_alignment(const gemmi::AlignmentResult& result,
                                const std::string& a, const std::string& b) {
  std::fputs(result.formatted(a, b).c_str(), stdout);
}

void print_result_summary(const gemmi::AlignmentResult& result) {
  printf("identity: %.2f%% / %.2f%%  CIGAR: %s\n",
         result.calculate_identity(1), result.calculate_identity(2),
         result.cigar_str().c_str());
}

std::vector<std::string> string_to_vector(const std::string& s) {
  std::vector<std::string> v(s.size());
  for (size_t i = 0; i != v.size(); ++i)
    v[i] = s[i];
  return v;
}

} // anonymous namespace

int GEMMI_MAIN(int argc, char **argv) {
  OptParser p(EXE_NAME);
  p.simple_parse(argc, argv, Usage);
  gemmi::AlignmentScoring scoring;
  if (p.options[Match])
    scoring.match = std::atoi(p.options[Match].arg);
  if (p.options[Mismatch])
    scoring.mismatch = std::atoi(p.options[Mismatch].arg);
  if (p.options[GapOpen])
    scoring.gapo = std::atoi(p.options[GapOpen].arg);
  if (p.options[GapExt])
    scoring.gape = std::atoi(p.options[GapExt].arg);
  bool verbose = p.options[Verbose];
  if ((bool)p.options[Query] != (bool)p.options[Target]) {
    std::fputs("Options --query and --target must be used together.\n", stderr);
    return 1;
  }
  p.check_exclusive_pair(TextAlign, Query);

  if (p.options[TextAlign]) {
    p.require_positional_args(2);
    if (p.nonOptionsCount() != 2) {
      std::fputs("two input strings are expected with --text-align\n", stderr);
      return 1;
    }
    std::string text1 = p.nonOption(0);
    std::string text2 = p.nonOption(1);
    std::vector<bool> free_gapo(1, 1);
    gemmi::AlignmentResult result
            = gemmi::align_string_sequences(string_to_vector(text1),
                                            string_to_vector(text2),
                                            free_gapo, scoring);
    printf("Score: %d   CIGAR: %s\n", result.score, result.cigar_str().c_str());
    if (p.options[PrintOneLetter])
      print_one_letter_alignment(result, text1, text2);
    return 0;
  }

  try {
    if (p.options[Query]) {
      int n_files = p.nonOptionsCount();
      if (n_files != 2 && n_files != 1)
        gemmi::fail("one or two input files are expected with --query/--target");
      using gemmi::read_structure_gz;
      std::vector<std::string> query;
      gemmi::PolymerType ptype = gemmi::PolymerType::Unknown;
      gemmi::Structure st1 = read_structure_gz(p.coordinate_input_file(0));
      if (p.options[Query].arg[0] == '+') {
        const gemmi::Entity* ent = get_entity(st1, p.options[Query].arg + 1);
        query = ent->full_sequence;
        ptype = ent->polymer_type;
      } else {
        query = get_polymer(get_first_model(st1), p.options[Query].arg)
                .extract_sequence();
      }
      gemmi::AlignmentResult result;
      gemmi::Structure st2_;
      if (n_files == 2)
        st2_ = read_structure_gz(p.coordinate_input_file(1));
      gemmi::Structure& st2 = n_files == 2 ? st2_ : st1;
      if (p.options[Target].arg[0] == '+') {
        const gemmi::Entity* ent = get_entity(st2, p.options[Target].arg + 1);
        std::vector<bool> free_gapo(1, 1);
        result = gemmi::align_string_sequences(query, ent->full_sequence,
                                               free_gapo, scoring);
        print_result_summary(result);
        if (p.options[PrintOneLetter])
          print_one_letter_alignment(result, gemmi::one_letter_code(query),
                                   gemmi::one_letter_code(ent->full_sequence));
      } else {
        auto polymer = get_polymer(get_first_model(st2), p.options[Target].arg);
        if (ptype == gemmi::PolymerType::Unknown)
          if (const gemmi::Entity* ent = st2.get_entity_of(polymer))
            ptype = ent->polymer_type;
        result = gemmi::align_sequence_to_polymer(query, polymer, ptype, scoring);
        print_result_summary(result);
        if (p.options[PrintOneLetter])
          print_one_letter_alignment(result, gemmi::one_letter_code(query),
                                     gemmi::one_letter_code(polymer));
      }
      if (p.options[Rmsd]) {
        auto poly1 = get_polymer(st1.models.at(0), p.options[Query].arg);
        auto poly2 = get_polymer(st2.models.at(0), p.options[Target].arg);
        printf("We superpose polymers using only matching residues and atoms w/o altloc.\n");
        gemmi::SupResult r;
        r = gemmi::calculate_superposition(poly1, poly2, ptype, gemmi::SupSelect::All);
        printf("RMSD of %zu atoms: %g\n", r.count, r.rmsd);
        r = gemmi::calculate_superposition(poly1, poly2, ptype, gemmi::SupSelect::CaP);
        printf("RMSD of %zu CA/P atoms: %g\n", r.count, r.rmsd);

        // this last part is not particularly useful
        auto mpoly2 = st2.models[0].find_chain(p.options[Target].arg)->get_polymer();
        for (gemmi::Residue& res : mpoly2)
          for (gemmi::Atom& atom : res.atoms)
            atom.pos = gemmi::Position(r.transform.apply(atom.pos));
        r = gemmi::calculate_current_rmsd(poly1, mpoly2, ptype, gemmi::SupSelect::All);
        printf("   the same rotation+shift applied to %zu atoms: %g\n", r.count, r.rmsd);
      }
      return 0;
    }

    p.require_input_files_as_args();
    for (int i = 0; i < p.nonOptionsCount(); ++i) {
      std::string input = p.coordinate_input_file(i);
      if (i > 0)
        printf("\n");
      if (verbose || p.nonOptionsCount() > 1)
        printf("File: %s\n", input.c_str());
      gemmi::Structure st = gemmi::read_structure_gz(input);
      const gemmi::Model& model = get_first_model(st);
      if (st.models.size() > 1)
        printf("Warning: using only model 1 of %zu.\n", st.models.size());
      for (const gemmi::Chain& chain : model.chains) {
        gemmi::ConstResidueSpan polymer = chain.get_polymer();
        if (!polymer)
          continue;
        const gemmi::Entity* ent = st.get_entity_of(polymer);
        if (!ent)
          gemmi::fail("No sequence (SEQRES) for chain " + chain.name);
        if (gemmi::seqid_matches_seqres(polymer, *ent))
          printf("Sequence numbers are wrt the full sequence (SEQRES).\n");
        gemmi::AlignmentResult result =
            gemmi::align_sequence_to_polymer(ent->full_sequence, polymer,
                                             ent->polymer_type, scoring);
        printf("%s chain %s  ", st.name.c_str(), chain.name.c_str());
        print_result_summary(result);
        if (p.options[CheckMmcif]) {
          if (st.input_format == gemmi::CoorFormat::Pdb)
            printf("Option --check-mmcif ignored for PDB file: %s\n", input.c_str());
          else
            check_label_seq_id(result, polymer);
        }
        if (p.options[PrintOneLetter])
          print_one_letter_alignment(result,
                                     gemmi::one_letter_code(ent->full_sequence),
                                     gemmi::one_letter_code(polymer));
        if (verbose)
          print_alignment_details(result, chain.name, polymer, *ent);
      }
    }
  } catch (std::runtime_error& e) {
    std::fprintf(stderr, "ERROR: %s\n", e.what());
    return 1;
  }
  return 0;
}
