package net.sf.okapi.steps.heuristicaligner;
   
   import java.util.ArrayList;
   import java.util.List;
   
   import com.acumenvelocity.ath.common.OkapiUtil;
   import com.acumenvelocity.ath.steps.BaseAlignerStep;
   
   import net.sf.okapi.common.IParameters;
  import net.sf.okapi.common.IResource;
  import net.sf.okapi.common.UsingParameters;
  import net.sf.okapi.common.filters.IFilter;
  import net.sf.okapi.common.resource.AlignmentStatus;
  import net.sf.okapi.common.resource.ITextUnit;
  import net.sf.okapi.common.resource.Segment;
  import net.sf.okapi.common.resource.TextContainer;
  import net.sf.okapi.common.resource.TextFragment;
  import net.sf.okapi.common.resource.TextPart;
  
  /**
   * Final, production-ready heuristic sentence aligner.
   * Supports M:N paragraph matching with sophisticated heuristics.
   */
  @UsingParameters(HeuristicSentenceAlignerParameters.class)
  public class HeuristicSentenceAlignerStep extends BaseAlignerStep {
  
    private HeuristicSentenceAlignerParameters params;
    private HeuristicAligner aligner;
  
    public HeuristicSentenceAlignerStep(IFilter targetFilter) {
      super(targetFilter);
      params = new HeuristicSentenceAlignerParameters();
      aligner = new HeuristicAligner();
    }
  
    @Override
    public String getName() {
      return "Improved Heuristic Sentence Alignment";
    }
  
    @Override
    public String getDescription() {
      return "Aligns paragraphs (M:N) and sentences using DP + translation-aware re-segmentation. "
          + "Uses sophisticated heuristics with back-translation for both levels.";
    }
  
    @Override
    public HeuristicSentenceAlignerParameters getParameters() {
      return params;
    }
  
    @Override
    public void setParameters(IParameters params) {
      this.params = (HeuristicSentenceAlignerParameters) params;
    }
  
    @Override
    protected boolean isSegmentSource() {
      return params.isSegmentSource();
    }
  
    @Override
    protected boolean isSegmentTarget() {
      return params.isSegmentTarget();
    }
  
    @Override
    protected boolean isUseCustomSourceRules() {
      return params.isUseCustomSourceRules();
    }
  
    @Override
    protected boolean isUseCustomTargetRules() {
      return params.isUseCustomTargetRules();
    }
  
    @Override
    protected String getCustomSourceRulesPath() {
      return params.getCustomSourceRulesPath();
    }
  
    @Override
    protected String getCustomTargetRulesPath() {
      return params.getCustomTargetRulesPath();
    }
  
    @Override
    protected boolean isCollapseWhitespace() {
      return params.isCollapseWhitespace();
    }
  
    @Override
    protected void performAlignment(List<ITextUnit> sourceTUs, List<ITextUnit> targetTUs) {
      filterMtArtifacts(sourceTUs, targetTUs);
  
      LOGGER.info("Alignment: {} source, {} target paragraphs", sourceTUs.size(), targetTUs.size());
  
      aligner.setTranslationAwareResegmentation(true);
  
     if (params.isForceSimpleOneToOneAlignment() && sourceTUs.size() == targetTUs.size()
         && sourceTUs.size() > 0) {
       LOGGER.info("Forcing 1:1 alignment");
       performSimpleOneToOneAlignment(sourceTUs, targetTUs);
 
     } else {
       // Use M:N paragraph matching with DP
       performManyToManyAlignment(sourceTUs, targetTUs);
     }
 
     // Set alignment origin on all remaining source TUs
     for (ITextUnit tu : sourceTUs) {
       OkapiUtil.setAlOrigin(tu, getSourceLocale(), getTargetLocale());
     }
   }
 
   /**
    * Simple 1:1 alignment when paragraph counts match
    */
   private void performSimpleOneToOneAlignment(List<ITextUnit> sourceTUs,
       List<ITextUnit> targetTUs) {
 
     for (int i = 0; i < sourceTUs.size(); i++) {
       ITextUnit srcTu = sourceTUs.get(i);
       ITextUnit trgTu = targetTUs.get(i);
 
       alignSentencesInParagraphPair(srcTu, trgTu);
 
       LOGGER.debug("1:1 aligned paragraph pair {}", i);
     }
   }
 
   /**
    * M:N paragraph alignment using dynamic programming.
    * Finds optimal grouping of source and target paragraphs.
    * Marks consumed TUs as non-translatable so they're filtered by base class.
    */
   private void performManyToManyAlignment(List<ITextUnit> sourceTUs, List<ITextUnit> targetTUs) {
     int m = sourceTUs.size();
     int n = targetTUs.size();
 
     // Build similarity matrix between all source and target paragraphs
     double[][] simMatrix = new double[m][n];
 
     for (int i = 0; i < m; i++) {
       String srcText = getPlainText(sourceTUs.get(i));
 
       for (int j = 0; j < n; j++) {
         String trgText = getPlainText(targetTUs.get(j));
         simMatrix[i][j] = aligner.calculateParagraphSimilarity(
             srcText.toLowerCase().trim(),
             trgText.toLowerCase().trim(),
             getSourceLocale(),
             getTargetLocale());
       }
     }
 
     // Find paragraph groups using a windowed approach
     List<ParagraphMatch> matches = findParagraphMatches(simMatrix);
 
     // Apply the matches - consumed TUs are marked as non-translatable
     int totalConsumed = 0;
     for (ParagraphMatch match : matches) {
       List<ITextUnit> consumed = applyParagraphMatch(match, sourceTUs, targetTUs);
       totalConsumed += consumed.size();
     }
 
     if (totalConsumed > 0) {
       LOGGER.info("Marked {} source TUs as non-translatable (merged into other TUs)",
           totalConsumed);
     }
   }
 
   /**
    * Find paragraph matches using a greedy windowed approach.
    * Groups consecutive source paragraphs and matches them to target groups.
    */
   private List<ParagraphMatch> findParagraphMatches(double[][] simMatrix) {
     List<ParagraphMatch> matches = new ArrayList<>();
     int m = simMatrix.length;
     int n = simMatrix[0].length;
 
     boolean[] srcUsed = new boolean[m];
     boolean[] trgUsed = new boolean[n];
 
     int srcPtr = 0;
     int trgPtr = 0;
 
     while (srcPtr < m) {
       // Try different source window sizes (1 to 6 paragraphs)
       ParagraphMatch bestMatch = null;
       double bestScore = -1.0;
 
       for (int srcWindow = 1; srcWindow <= Math.min(6, m - srcPtr); srcWindow++) {
         // Check if any source in this window is already used
         boolean srcConflict = false;
         for (int s = srcPtr; s < srcPtr + srcWindow; s++) {
           if (srcUsed[s]) {
             srcConflict = true;
             break;
           }
         }
         if (srcConflict)
           break;
 
         // Try different target window sizes
         for (int trgWindow = 0; trgWindow <= Math.min(6, n - trgPtr); trgWindow++) {
           // Check if any target in this window is already used
           boolean trgConflict = false;
           for (int t = trgPtr; t < trgPtr + trgWindow; t++) {
             if (trgUsed[t]) {
               trgConflict = true;
               break;
             }
           }
           if (trgConflict)
             break;
 
           // Calculate combined similarity for this match
           double score = calculateGroupSimilarity(
               simMatrix, srcPtr, srcWindow, trgPtr, trgWindow);
 
           // Penalize mismatched group sizes
           double sizePenalty = Math.abs(srcWindow - trgWindow) * 0.05;
           double adjustedScore = score - sizePenalty;
 
           if (adjustedScore > bestScore) {
             bestScore = score; // Store unadjusted for threshold
             bestMatch = new ParagraphMatch(srcPtr, srcWindow, trgPtr, trgWindow, score);
           }
         }
       }
 
       if (bestMatch != null && (bestMatch.targetCount > 0
           ? bestScore >= HeuristicAligner.PARAGRAPH_MATCH_THRESHOLD
           : bestMatch.sourceCount <= 3)) { // Allow small deletions
 
         // Mark as used
         for (int s = bestMatch.sourceStart; s < bestMatch.sourceStart
             + bestMatch.sourceCount; s++) {
           srcUsed[s] = true;
         }
         for (int t = bestMatch.targetStart; t < bestMatch.targetStart
             + bestMatch.targetCount; t++) {
           trgUsed[t] = true;
         }
 
         matches.add(bestMatch);
         srcPtr = bestMatch.sourceStart + bestMatch.sourceCount;
         trgPtr = bestMatch.targetStart + bestMatch.targetCount;
 
         LOGGER.info("Paragraph match: s[{}..{}] -> t[{}..{}] (score: {:.3f})",
             bestMatch.sourceStart, bestMatch.sourceStart + bestMatch.sourceCount - 1,
             bestMatch.targetStart, bestMatch.targetStart + bestMatch.targetCount - 1,
             bestScore);
       } else {
         // No good match, advance by 1
         srcPtr++;
       }
     }
 
     return matches;
   }
 
   /**
    * Calculate average similarity between a group of source and target paragraphs
    */
   private double calculateGroupSimilarity(double[][] simMatrix,
       int srcStart, int srcCount, int trgStart, int trgCount) {
 
     if (srcCount == 0 || trgCount == 0) {
       return 0.0; // Deletion or insertion
     }
 
     double totalScore = 0.0;
     int count = 0;
 
     for (int i = srcStart; i < srcStart + srcCount; i++) {
       for (int j = trgStart; j < trgStart + trgCount; j++) {
         totalScore += simMatrix[i][j];
         count++;
       }
     }
 
     return count > 0 ? totalScore / count : 0.0;
   }
 
   /**
    * Apply a paragraph match by merging source and target paragraphs.
    * Returns list of consumed source TUs that should be marked as non-translatable.
    */
   private List<ITextUnit> applyParagraphMatch(ParagraphMatch match, List<ITextUnit> sourceTUs,
       List<ITextUnit> targetTUs) {
 
     List<ITextUnit> consumedTUs = new ArrayList<>();
 
     // Collect source TUs
     List<ITextUnit> srcGroup = new ArrayList<>();
     for (int i = match.sourceStart; i < match.sourceStart + match.sourceCount; i++) {
       srcGroup.add(sourceTUs.get(i));
     }
 
     // Collect target TUs
     List<ITextUnit> trgGroup = new ArrayList<>();
     for (int i = match.targetStart; i < match.targetStart + match.targetCount; i++) {
       trgGroup.add(targetTUs.get(i));
     }
 
     if (srcGroup.isEmpty()) {
       return consumedTUs;
     }
 
     // If srcGroup and trgGroup contain the same number of TUs, we align the TUs 1:1
     if (srcGroup.size() == trgGroup.size()) {
       for (int i = 0; i < srcGroup.size(); i++) {
         ITextUnit srcTu = srcGroup.get(i);
         ITextUnit trgTu = trgGroup.get(i);
         
         alignSentencesInParagraphPair(srcTu, trgTu);
       }
       
       return consumedTUs;
     }
 
     // Primary source TU (first in group) - this is the one we'll keep
     ITextUnit primarySrc = srcGroup.get(0);
 
     // Merge additional source TUs into primary
     for (int i = 1; i < srcGroup.size(); i++) {
       ITextUnit additionalSrc = srcGroup.get(i);
 
       LOGGER.debug("Merging source TU {} into {}",
           additionalSrc.getId(), primarySrc.getId());
 
       // Merge all segments from additional TU into primary
       for (Segment seg : additionalSrc.getSource().getSegments()) {
         primarySrc.getSource().append(seg.clone());
       }
 
       // Mark as non-translatable and clear source so it won't be output
       additionalSrc.setIsTranslatable(false);
       additionalSrc.getSource().clear();
 
       consumedTUs.add(additionalSrc);
     }
 
     if (trgGroup.isEmpty()) {
       // No target - create empty
       LOGGER.debug("Creating empty target for source group starting at {}", match.sourceStart);
       createEmptyTargetSegments(primarySrc);
 
     } else {
       // Merge target TUs into a single virtual target
       ITextUnit mergedTarget = trgGroup.get(0).clone();
 
       for (int i = 1; i < trgGroup.size(); i++) {
         ITextUnit additionalTrg = trgGroup.get(i);
 
         for (Segment seg : additionalTrg.getSource().getSegments()) {
           mergedTarget.getSource().append(seg.clone());
         }
       }
 
       // Align sentences between merged source and merged target
       alignSentencesInParagraphPair(primarySrc, mergedTarget);
     }
 
     return consumedTUs;
   }
 
   /**
    * Align sentences within a single paragraph pair using the heuristic aligner.
    */
   private void alignSentencesInParagraphPair(ITextUnit sourceTu, ITextUnit targetTu) {
     // Prepare source and target segment lists BEFORE modifying anything
     List<Segment> sourceSegments = new ArrayList<>(sourceTu.getSource().getSegments().asList());
     List<Segment> targetSegments = new ArrayList<>(targetTu.getSource().getSegments().asList());
 
     if (sourceSegments.isEmpty()) {
       LOGGER.warn("Source has no segments in paragraph pair");
       return;
     }
 
     if (targetSegments.isEmpty()) {
       LOGGER.warn("Target has no segments in paragraph pair");
       createEmptyTargetSegments(sourceTu);
       return;
     }
 
     // Temporarily set target so aligner can work with it
     TextContainer tempTarget = new TextContainer();
 
     for (Segment seg : targetSegments) {
       tempTarget.getSegments().append(seg.clone());
     }
 
     sourceTu.setTarget(getTargetLocale(), tempTarget);
 
     // Call aligner to get sentence matches
     List<SentenceMatch> matches = aligner.alignSentencesInTu(
         sourceTu, getSourceLocale(), getTargetLocale());
 
     // Create final target container and populate based on matches
     TextContainer finalTarget = sourceTu.createTarget(getTargetLocale(), false,
         IResource.CREATE_EMPTY);
 
     finalTarget.clear();
 
     // Apply alignment results
     for (SentenceMatch m : matches) {
       if (m.sourceIndex >= 0 && m.sourceIndex < sourceSegments.size()) {
         Segment srcSeg = sourceSegments.get(m.sourceIndex);
 
         if (m.targetIndex >= 0 && m.targetIndex < targetSegments.size()) {
           Segment trgSeg = targetSegments.get(m.targetIndex);
           finalTarget.append(new Segment(srcSeg.getId(), trgSeg.text.clone()));
 
         } else {
           // We got -1 for the tseg index matching the current sseg index
           if (sourceSegments.size() == 1 || targetSegments.size() == 1) {
             // 1:1, M:1, 1:N cases
             finalTarget = targetTu.getSource();
             sourceTu.setTarget(getTargetLocale(), finalTarget);
 
             if (targetSegments.size() == 1) {
               sourceTu.getSource().joinAll();
             }
 
             if (sourceSegments.size() == 1) {
               finalTarget.joinAll();
             }
 
             break;
 
           } else {
             // TODO test real M:N cases
             finalTarget.append(new Segment(srcSeg.getId(), new TextFragment("???")));
           }
         }
       }
     }
 
     finalTarget.setHasBeenSegmentedFlag(true);
     finalTarget.getSegments().setAlignmentStatus(AlignmentStatus.ALIGNED);
   }
 
   /**
    * Create empty target segments for a source TU with no match
    */
   private void createEmptyTargetSegments(ITextUnit tu) {
     tu.createTarget(getTargetLocale(), false, IResource.CREATE_EMPTY);
     for (Segment s : tu.getSource().getSegments()) {
       tu.getTarget(getTargetLocale()).append(new Segment(s.getId(), new TextFragment("")));
     }
   }
 
   /**
    * Filter out MT artifacts from target text units.
    * Detect and remove duplicate source paragraphs.
    * Filter source metadata but NOT target metadata.
    */
   private void filterMtArtifacts(List<ITextUnit> sourceTUs, List<ITextUnit> targetTUs) {
     // Filter targets - only MT artifacts, keep "passage" and other metadata
     List<ITextUnit> filtered = new ArrayList<>();
     for (ITextUnit tu : targetTUs) {
       String text = tu.getSource().toString().toLowerCase().trim();
 
       // Skip ONLY MT artifacts
       if (text.matches(".*\\bmachine\\b.*translated.*|.*google.*|.*auto.*translated.*")) {
         LOGGER.info("Filtered MT artifact: {}", text);
         continue;
       }
 
       // Keep everything else, including "passage" metadata
       filtered.add(tu);
     }
 
     int removedCount = targetTUs.size() - filtered.size();
 
     if (removedCount > 0) {
       LOGGER.info("Filtered {} MT artifact paragraphs from target", removedCount);
     }
 
     targetTUs.clear();
     targetTUs.addAll(filtered);
 
     // Filter source: metadata AND duplicates
     // List<ITextUnit> srcFiltered = new ArrayList<>();
     // List<String> seenTexts = new ArrayList<>();
     //
     // for (ITextUnit tu : sourceTUs) {
     // String text = tu.getSource().toString().toLowerCase().trim();
     //
     // // Skip metadata markers from source only
     // if (text.equals("passage") || text.isEmpty()) {
     // LOGGER.info("Filtered source metadata/empty: {}", text);
     // continue;
     // }
     //
     // // Check for duplicates
     // boolean isDuplicate = false;
     //
     // for (String seenText : seenTexts) {
     // if (text.equals(seenText)) {
     // LOGGER.info("Filtered duplicate source paragraph: {}",
     // text.length() > 50 ? text.substring(0, 50) + "..." : text);
     //
     // isDuplicate = true;
     // break;
     // }
     // }
     //
     // if (!isDuplicate) {
     // srcFiltered.add(tu);
     // seenTexts.add(text);
     // }
     // }
     //
     // if (sourceTUs.size() != srcFiltered.size()) {
     // LOGGER.info("Filtered {} source metadata/duplicate paragraphs",
     // sourceTUs.size() - srcFiltered.size());
     // sourceTUs.clear();
     // sourceTUs.addAll(srcFiltered);
     // }
   }
 
   /**
    * Extract plain text from a text unit
    */
   private String getPlainText(ITextUnit tu) {
     StringBuilder sb = new StringBuilder();
     for (TextPart tp : tu.getSource().getParts()) {
       if (sb.length() > 0) {
         sb.append(" ");
       }
       sb.append(tp.getContent().getText());
     }
     return sb.toString().trim();
   }
 
   /**
    * Data class for paragraph match results
    */
   static class ParagraphMatch {
     int sourceStart;
     int sourceCount;
     int targetStart;
     int targetCount;
     double score;
 
     ParagraphMatch(int ss, int sc, int ts, int tc, double score) {
       this.sourceStart = ss;
       this.sourceCount = sc;
       this.targetStart = ts;
       this.targetCount = tc;
       this.score = score;
     }
   }
 
   /**
    * Simple data class to hold sentence match results
    */
   static class SentenceMatch {
     public int sourceIndex;
     public int targetIndex = -1;
     public double score;
 
     public SentenceMatch(int s, int t, double sc) {
       this.sourceIndex = s;
       this.targetIndex = t;
       this.score = sc;
     }
   }
 }