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.model.ClinicalCategory;
import gov.cms.grouper.snf.model.CognitiveLevel;
import gov.cms.grouper.snf.model.SnfDiagnosisCode;
import gov.cms.grouper.snf.model.reader.Rai300;
import gov.cms.grouper.snf.model.table.BasicRow;
import gov.cms.grouper.snf.model.table.ClinicalCategoryMappingRow;
import gov.cms.grouper.snf.model.table.DiagnosisMasterRow;
import gov.cms.grouper.snf.model.table.SlpCmgRow;
import gov.cms.grouper.snf.model.table.SlpComorbiditiesRow;
import gov.cms.grouper.snf.util.ClaimInfo;
import java.util.List;
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=681" class="req">PDPM Payment
 * Component: SLP</a>
 */
public class SpeechLogic extends SnfDataVersionImpl<String> {

  private static final Logger log = LoggerFactory.getLogger(SpeechLogic.class);

  public static final Set<Rai300> slpRelatedComorbiditiesCheckBoxes = SnfUtils.toSet(Rai300.I4300,
      Rai300.I4500, Rai300.I4900, Rai300.I5500, Rai300.O0100E2, Rai300.O0100F2);

  public static final Set<Rai300> swallowingDisorders =
      SnfUtils.toSet(Rai300.K0100A, Rai300.K0100B, Rai300.K0100C, Rai300.K0100D);

  public static final Set<Rai300> presenceOfMechanicallyAlteredDiet =
      SnfUtils.toSet(Rai300.K0510C2);

  private final ClaimInfo claim;
  private final List<SnfDiagnosisCode> secondaryDiagnosisCodes;
  private final Supplier<Boolean> step6PresenceOfSwallowingDisorder;
  private final Supplier<Boolean> step7PresenceOfMechanicallyAlteredDiet;
  private final CognitiveLevel cognitiveLevel;
  private final ClinicalCategory pdxClinicalCategory;

  public SpeechLogic(int version, ClaimInfo claim, List<SnfDiagnosisCode> secondaryDiagnosisCodes,
      CognitiveLevel cognitiveLevel, ClinicalCategory pdxClinicalCategory) {
    this(version, claim, secondaryDiagnosisCodes, cognitiveLevel, pdxClinicalCategory, null, null);
  }

  protected SpeechLogic(int version, ClaimInfo claim,
      List<SnfDiagnosisCode> secondaryDiagnosisCodes, CognitiveLevel cognitiveLevel,
      ClinicalCategory pdxClinicalCategory, Supplier<Boolean> step6PresenceOfSwallowingDisorder,
      Supplier<Boolean> step7PresenceOfMechanicallyAlteredDiet) {
    super(version);
    this.claim = claim;
    this.secondaryDiagnosisCodes = secondaryDiagnosisCodes;
    this.cognitiveLevel = cognitiveLevel;
    this.pdxClinicalCategory = pdxClinicalCategory;

    this.step6PresenceOfSwallowingDisorder = SnfUtils.nullCheck(step6PresenceOfSwallowingDisorder,
        () -> claim.isAnyAssessmentValuesPresent(SpeechLogic.swallowingDisorders),
        step6PresenceOfSwallowingDisorder);

    this.step7PresenceOfMechanicallyAlteredDiet =
        SnfUtils.nullCheck(step7PresenceOfMechanicallyAlteredDiet,
            () -> claim.isAnyAssessmentValuesPresent(SpeechLogic.presenceOfMechanicallyAlteredDiet),
            step7PresenceOfMechanicallyAlteredDiet);
  }

  /**
   * Map the Primary Diagnosis (pdx) clinical category to a SLP related clinical category.
   * <a href="doc-files/mds-3.0-rai-manual-v1.17.1_october_2019.pdf#page=683" class="req">Step2</a>
   *
   * @return the mapped SLP clinical category based on the pdx clinical category
   */
  public static String step2SlpClinicalCategory(int version, String pdxClinicalCategory) {
    ClinicalCategoryMappingRow row = SnfTables.get(SnfTables.clinicalCategoryMapping,
        pdxClinicalCategory, BasicRow.getVersionSelector(), version);

    String result = row.getSlpClinicalCategory();
    return SnfContext.trace(result);
  }

  /**
   * Determine whether the resident has one or more SLP-related comorbidities
   * <a href="doc-files/mds-3.0-rai-manual-v1.17.1_october_2019.pdf#page=683" class="req">Step3</a>
   */
  public static boolean step3HasSlpRelatedComorbidities(int version,
      Supplier<Boolean> slpAssessmentsSupplier, List<SnfDiagnosisCode> secondaryDiagnoses) {

    boolean claimChecked = slpAssessmentsSupplier.get();

    boolean hasSlpRelatedComorbidities = false;

    final Set<SlpComorbiditiesRow> i8000SlpRows = SnfTables.getAll(SnfTables.slpComorbiditiesTable,
        SlpComorbiditiesRow.I8000, BasicRow.getVersionSelector(), version);
    final List<String> generalDiagnosisSlpConditions =
        i8000SlpRows.stream().map(SlpComorbiditiesRow::getCategory).collect(Collectors.toList());

    for (SnfDiagnosisCode code : secondaryDiagnoses) {
      // Look up dx code to determine its slp category
      DiagnosisMasterRow row = SnfTables.get(SnfTables.diagnosisMasterTable, code.getValue(),
          BasicRow.getVersionSelector(), version);
      String slpCategory = "";

      if (row == null) {
        log.warn("Invalid ICD 10 code: " + code);
      } else if (row.getSlpCategory() != null) {
        slpCategory = row.getSlpCategory();
      }

      // Check if diagnosis' slp category is related
      if (generalDiagnosisSlpConditions.contains(slpCategory)) {
        hasSlpRelatedComorbidities = true;
        break;
      }
    }

    boolean result = claimChecked || hasSlpRelatedComorbidities;
    return SnfContext.trace(result);
  }

  /**
   * Determine whether the resident has a cognitive impairment
   * <a href="doc-files/mds-3.0-rai-manual-v1.17.1_october_2019.pdf#page=684" class="req">Step4</a>
   */
  public static boolean step4HasCognitiveImpairment(CognitiveLevel cognitiveLevel) {
    return CognitiveLevel.COGNITIVELY_INTACT != cognitiveLevel;
  }

  /**
   * Determine how many of the following conditions are present
   * <a href="doc-files/mds-3.0-rai-manual-v1.17.1_october_2019.pdf#page=684" class="req">Step5</a>
   */
  public static int step5NumberOfConditionsPresent(
      Supplier<String> step2SlpClinicalCategorySupplier,
      Supplier<Boolean> step3HasSlpRelatedComorbiditiesSupplier,
      Supplier<Boolean> step4HasCognitiveImpairmentSupplier) {

    int numberOfConditionsPresent = ClinicalCategoryMappingRow.SLP_Acute_Neurologic
        .equals(step2SlpClinicalCategorySupplier.get()) ? 1 : 0;

    numberOfConditionsPresent += step3HasSlpRelatedComorbiditiesSupplier.get() ? 1 : 0;

    numberOfConditionsPresent += step4HasCognitiveImpairmentSupplier.get() ? 1 : 0;

    int result = numberOfConditionsPresent;
    return SnfContext.trace(result);
  }

  public static final int NEITHER = 0;
  public static final int EITHER = 1;
  public static final int BOTH = 2;

  /**
   * <a href="doc-files/mds-3.0-rai-manual-v1.17.1_october_2019.pdf#page=684" class="req">Step
   * 6</a>: Determine whether the resident has a swallowing disorder using item K0100
   * <a href="doc-files/mds-3.0-rai-manual-v1.17.1_october_2019.pdf#page=684" class="req">Step
   * 7</a>: Determine whether the resident has a mechanically altered diet
   * <a href="doc-files/mds-3.0-rai-manual-v1.17.1_october_2019.pdf#page=684" class="req">Step
   * 8</a>: Determine how many of the following conditions are present based on Steps 6 and 7
   */
  public static int step6To8PresenceOfMechanicallyAlteredDietOrSwallowingDisorder(
      Supplier<Boolean> step6PresenceOfSwallowingDisorder,
      Supplier<Boolean> step7PresenceOfMechanicallyAlteredDiet) {
    boolean step6IsSwallowingDisorder = step6PresenceOfSwallowingDisorder.get();
    boolean step7IsMechanicallyAlteredDiet = step7PresenceOfMechanicallyAlteredDiet.get();
    int result = SpeechLogic.EITHER; // Has either

    if (step6IsSwallowingDisorder && step7IsMechanicallyAlteredDiet) {
      result = SpeechLogic.BOTH; // Has both
    } else if (!step6IsSwallowingDisorder && !step7IsMechanicallyAlteredDiet) {
      result = SpeechLogic.NEITHER; // Has neither
    }

    return SnfContext.trace(result);
  }

  /**
   * Determine the resident’s SLP case-mix group using the responses from Steps 1-8
   * <a href="doc-files/mds-3.0-rai-manual-v1.17.1_october_2019.pdf#page=685" class="req">Step9</a>
   */
  public static SlpCmgRow step9PdpmSlpClassification(int version,
      Supplier<Integer> step5NumberOfConditionsPresent,
      Supplier<Integer> step6To8PresenceOfMechanicallyAlteredDietOrSwallowingDisorder) {

    Integer conditionsPresentCount = step5NumberOfConditionsPresent.get();
    // 0 - Has neither, 1 - Has either, 2 - Has both
    Integer mechAltDietOrSwallowDisorder =
        step6To8PresenceOfMechanicallyAlteredDietOrSwallowingDisorder.get();

    Set<SlpCmgRow> rows = SnfTables.getAll(SnfTables.slpCmgTable, conditionsPresentCount,
        BasicRow.getVersionSelector(), version);

    // Filter further on Altered Diet or Swallowing Disorder
    SlpCmgRow row = SnfUtils.doOrDie(() -> rows.stream().filter((item) -> {
      boolean result = mechAltDietOrSwallowDisorder.intValue() == item
          .getMechAltDietOrSwallowDisorder().intValue();
      return result;
    }).findFirst().get());

    return row;
  }

  @Override
  public String exec() {
    Supplier<String> step2SlpClinicalCategory = () -> SpeechLogic
        .step2SlpClinicalCategory(this.getDataVersion(), this.pdxClinicalCategory.getDescription());

    Supplier<Boolean> step3HasSlpRelatedComorbidities =
        () -> SpeechLogic.step3HasSlpRelatedComorbidities(this.getDataVersion(),
            () -> this.claim
                .isAnyAssessmentValuesPresent(SpeechLogic.slpRelatedComorbiditiesCheckBoxes),
            this.secondaryDiagnosisCodes);

    Supplier<Boolean> step4HasCognitiveImpairment =
        () -> SpeechLogic.step4HasCognitiveImpairment(this.cognitiveLevel);

    Supplier<Integer> step5NumberOfConditionsPresent =
        () -> SpeechLogic.step5NumberOfConditionsPresent(step2SlpClinicalCategory,
            step3HasSlpRelatedComorbidities, step4HasCognitiveImpairment);

    Supplier<Integer> step6To8PresenceOfMechanicallyAlteredDietOrSwallowingDisorder =
        () -> SpeechLogic.step6To8PresenceOfMechanicallyAlteredDietOrSwallowingDisorder(
            this.step6PresenceOfSwallowingDisorder, this.step7PresenceOfMechanicallyAlteredDiet);

    String cmg = SpeechLogic
        .step9PdpmSlpClassification(this.getDataVersion(), step5NumberOfConditionsPresent,
            step6To8PresenceOfMechanicallyAlteredDietOrSwallowingDisorder)
        .getCmg();

    return SnfContext.trace(cmg);
  }

}
