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

import gov.cms.grouper.snf.SnfContext;
import gov.cms.grouper.snf.SnfTables;
import gov.cms.grouper.snf.model.Assessment;
import gov.cms.grouper.snf.model.SnfProcessException;
import gov.cms.grouper.snf.model.enums.AssessmentType;
import gov.cms.grouper.snf.model.enums.ClinicalCategory;
import gov.cms.grouper.snf.model.enums.PdxEligibility;
import gov.cms.grouper.snf.model.enums.Rai300;
import gov.cms.grouper.snf.model.table.BasicRow;
import gov.cms.grouper.snf.model.table.ClinicalCategoryMasterRow;
import gov.cms.grouper.snf.model.table.CmgForPtOtRow;
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.BiFunction;
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> {
  private static final Logger log = LoggerFactory.getLogger(TherapyLogic.class);
  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 final BiFunction<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;


  /**
   * Physical Therapy logic is the same as Occupational Therapy
   */
  public TherapyLogic(int version, ClaimInfo claim, ClinicalCategory pdxClinicalCategory) {
    super(version);
    this.pdxClinicalCategory = pdxClinicalCategory;
    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));
    functionalAssessmentsMap.forEach((item, score) -> functionalAssessmentsMap.put(item,
        claim.performanceRecode(() -> score)));
    functionalAssessments = (assessmentList, code) ->
        functionalAssessmentsMap.get(code) == null ? (Integer) 0 : functionalAssessmentsMap.get(code);
    step3Supplier = () -> claim.calculateFunctionScoreString(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) {

    ClinicalCategoryMasterRow clinicalCategoryMasterRow = SnfTables
        .get(SnfTables.clinicalCategoryMasterTable, pdx, BasicRow::isVersion, version);

    if (clinicalCategoryMasterRow == null) {
      log.warn("Invalid ICD 10 code: " + pdx);
      return SnfContext.trace("Null clinical category because of invalid PDX " + pdx, null);
    }

    PdxEligibility pdxEligibility = clinicalCategoryMasterRow.getPdxEligibility();

    ClinicalCategory result = ClinicalCategory.of(clinicalCategoryMasterRow.getClinicalCategory())
            .orElseThrow(() -> new SnfProcessException(
                    "ClinicalCategory " + clinicalCategoryMasterRow.getClinicalCategory()
                            + " not recognized. Data issue!"));

    if (pdxEligibility == PdxEligibility.ORTHOPEDIC_SURGERY) {
      if (asstInfo.isCheckedAndNotNull(Rai300.J2100)) {
        // Step #1A/B
        if (asstInfo.isAnyAssessmentValuesPresent(
            Rai300.getAssessmentsByType(AssessmentType.PROCEDURE_MAJOR), "MAJOR PROCEDURES")) {
          result = ClinicalCategory.MAJOR_JOINT_REPLACEMENT_OR_SPINAL_SURGERY;
        } else if (asstInfo.isAnyAssessmentValuesPresent(
            Rai300.getAssessmentsByType(AssessmentType.PROCEDURE_ORTHOPEDIC), "ORTHOPEDIC PROCEDURES")) {
          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), "NON-ORTHOPEDIC PROCEDURES")) {
          result = ClinicalCategory.NON_ORTHOPEDIC_SURGERY;
        }
      }
    }

    return SnfContext.trace(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) {
    String clinicalCategory = SnfTables.get(SnfTables.clinicalCategoryMapping, pdxClinicalCategory.getDescription(),
            BasicRow::isVersion, version).getPtOtClinicalCategory();
    ClinicalCategory result = ClinicalCategory
        .of(clinicalCategory)
            .orElseThrow(() -> new SnfProcessException(
                    "ClinicalCategory " + clinicalCategory + " not recognized. Data issue!"));
    return SnfContext.trace("Determined clinical category", 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::isVersion, version);

    Integer determinedScore = functionScoreSupplier.get();
    CmgForPtOtRow row = list.stream()
            .filter(item -> item.getFunctionScoreLow() <= determinedScore
              && item.getFunctionScoreHigh() >= determinedScore)
            .findFirst().orElse(null);

    if (row == null) {
      throw new SnfProcessException("Unable to determine CMG for Therapy Logic");
    }

    return SnfContext.trace(row);
  }

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

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

    return SnfContext.trace("----------- Therapy CMG", result.getCmg());
  }
}
