package gov.cms.grouper.snf.component.r2.logic;

import gov.cms.grouper.snf.SnfContext;
import gov.cms.grouper.snf.SnfTables;
import gov.cms.grouper.snf.lego.SnfUtils;
import gov.cms.grouper.snf.lego.SnfDataVersionImpl;
import gov.cms.grouper.snf.lego.TriFunction;
import gov.cms.grouper.snf.model.Assessment;
import gov.cms.grouper.snf.model.AssessmentType;
import gov.cms.grouper.snf.model.ClinicalCategory;
import gov.cms.grouper.snf.model.PdxEligibility;
import gov.cms.grouper.snf.model.reader.Rai300;
import gov.cms.grouper.snf.model.table.BasicRow;
import gov.cms.grouper.snf.model.table.CmgForPtOtRow;
import gov.cms.grouper.snf.model.table.DiagnosisMasterRow;
import gov.cms.grouper.snf.util.ClaimInfo;
import java.util.Arrays;
import java.util.List;
import java.util.Map;
import java.util.Set;
import java.util.function.Supplier;
import java.util.stream.Collectors;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

/**
 * <a href="doc-files/mds-3.0-rai-manual-v1.17.1_october_2019.pdf#page=671" class="req">PDPM Payment
 * Component: PT</a>
 */
public class TherapyLogic extends SnfDataVersionImpl<String> {

  protected static final List<String> bedMobilityList = Arrays.asList("GG0170B", "GG0170C");
  protected static final List<String> transferList = Arrays.asList("GG0170D", "GG0170E", "GG0170F");
  protected static final List<String> walkingList = Arrays.asList("GG0170J", "GG0170K");
  protected static final List<String> generalItemList =
      Arrays.asList("GG0130A", "GG0130B", "GG0130C");
  private static final Logger log = LoggerFactory.getLogger(TherapyLogic.class);
  private final ClaimInfo claim;
  private final TriFunction<Integer, ClaimInfo, String, Integer> functionalAssessments;
  private final Map<String, Integer> functionalAssessmentsMap;
  private final ClinicalCategory pdxClinicalCategory;
  /**
   * Calculate the resident's function score for PT/OT payments
   */
  private final Supplier<Integer> step3Supplier;


  public TherapyLogic(int version, ClaimInfo claim, ClinicalCategory pdxClinicalCategory) {
    super(version);
    this.pdxClinicalCategory = pdxClinicalCategory;
    this.claim = claim;
    this.functionalAssessmentsMap =
        claim.hasIpa()
            ? claim.getAssessments().stream()
                .filter(item -> item.getItem().startsWith("GG") && item.getItem().endsWith("5"))
                .collect(Collectors.toMap(
                    item -> item.getItem().substring(0, item.getItem().length() - 1),
                    Assessment::getValueInt))
            : claim.getAssessments().stream()
                .filter(item -> item.getItem().startsWith("GG") && item.getItem().endsWith("1"))
                .collect(Collectors.toMap(
                    item -> item.getItem().substring(0, item.getItem().length() - 1),
                    Assessment::getValueInt));
    this.functionalAssessmentsMap.forEach((item, score) -> this.functionalAssessmentsMap.put(item,
        claim.performanceRecode(() -> score)));
    this.functionalAssessments =
        (ver, assessmentList, code) -> SnfUtils.nullCheck(this.functionalAssessmentsMap.get(code),
            0, this.functionalAssessmentsMap.get(code));
    this.step3Supplier = () -> claim.calculateFunctionScoreString(this.functionalAssessments,
        TherapyLogic.bedMobilityList, TherapyLogic.transferList, TherapyLogic.walkingList,
        TherapyLogic.generalItemList);
  }


  /**
   * Go thru a set of logic to find the true primary diagnosis, if not found then use PDX's Clinical
   * Category
   * <a href="doc-files/mds-3.0-rai-manual-v1.17.1_october_2019.pdf#page=671" class="req">step1</a>
   *
   * @return primary clinical category
   */
  public static ClinicalCategory getPdxClinicalCategory(int version, ClaimInfo asstInfo,
      String pdx) {

    DiagnosisMasterRow diagnosisMasterRow =
        SnfTables.get(SnfTables.diagnosisMasterTable, pdx, BasicRow.getVersionSelector(), version);

    if (diagnosisMasterRow == null) {
      log.warn("Invalid ICD 10 code: " + pdx);
      return null;
    }

    PdxEligibility pdxEligibility = diagnosisMasterRow.getPdxEligibility();

    ClinicalCategory result = ClinicalCategory.of(diagnosisMasterRow.getClinicalCategory());

    if (pdxEligibility == PdxEligibility.ORTHOPEDIC_SURGERY) {
      if (asstInfo.isCheckedAndNotNull(Rai300.J2100)) {
        // Step #1A/B
        if (asstInfo.isAnyAssessmentValuesPresent(
            Rai300.getAssessmentsByType(AssessmentType.PROCEDURE_MAJOR))) {
          result = ClinicalCategory.MAJOR_JOINT_REPLACEMENT_OR_SPINAL_SURGERY;
        } else if (asstInfo.isAnyAssessmentValuesPresent(
            Rai300.getAssessmentsByType(AssessmentType.PROCEDURE_ORTHOPEDIC))) {
          result = ClinicalCategory.ORTHOPEDIC_SURGERY;
        }
      }
    } else if (pdxEligibility == PdxEligibility.NON_ORTHOPEDIC_SURGERY) {
      if (asstInfo.isCheckedAndNotNull(Rai300.J2100)) {
        if (asstInfo.isAnyAssessmentValuesPresent(
            Rai300.getAssessmentsByType(AssessmentType.PROCEDURE_NON_ORTHOPEDIC))) {
          result = ClinicalCategory.NON_ORTHOPEDIC_SURGERY;
        }
      }
    }

    return result;
  }

  /**
   * Determine the resident’s clinical category based on the mapping in
   * Clinical_Category_Mapping.csv
   * <a href="doc-files/mds-3.0-rai-manual-v1.17.1_october_2019.pdf#page=673" class="req">step2</a>
   *
   * @return Therapy's clinical category
   */
  protected static ClinicalCategory step2(int version, ClinicalCategory pdxClinicalCategory) {
    ClinicalCategory result = ClinicalCategory
        .of(SnfTables.get(SnfTables.clinicalCategoryMapping, pdxClinicalCategory.getDescription(),
            BasicRow.getVersionSelector(), version).getPtOtClinicalCategory());
    return SnfContext.trace(result);
  }

  /**
   * Using the responses from Steps 2 and 3 above, determine the resident’s case-mix group using the
   * table
   * <a href="doc-files/mds-3.0-rai-manual-v1.17.1_october_2019.pdf#page=675" class="req">step4</a>
   *
   * @return Case-mix group
   */
  protected static CmgForPtOtRow step4(int version,
      Supplier<ClinicalCategory> clinicalCategorySupplier,
      Supplier<Integer> functionScoreSupplier) {

    Set<CmgForPtOtRow> list = SnfTables.getAll(SnfTables.cmgForPtOt,
        clinicalCategorySupplier.get().getDescription(), BasicRow.getVersionSelector(), version);

    CmgForPtOtRow row = list.stream().filter(item -> {
      boolean result = item.getFunctionScoreLow() <= functionScoreSupplier.get()
          && item.getFunctionScoreHigh() >= functionScoreSupplier.get();
      return result;
    }).findFirst().orElse(null);

    return SnfContext.trace(SnfUtils.doOrDie(row::getCmg), row);
  }

  @Override
  public String exec() {
    Supplier<ClinicalCategory> clinicalCategory =
        () -> TherapyLogic.step2(this.getDataVersion(), this.pdxClinicalCategory);

    CmgForPtOtRow result =
        TherapyLogic.step4(this.getDataVersion(), clinicalCategory, this.step3Supplier);

    return SnfContext.trace(result.getCmg());
  }
}
