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.enums.CognitiveLevel;
import gov.cms.grouper.snf.model.enums.Rai300;
import gov.cms.grouper.snf.model.table.CognitiveLevelRow;
import gov.cms.grouper.snf.util.ClaimInfo;
import java.util.Collection;
import java.util.function.Supplier;

/**
 * <a href="doc-files/mds-3.0-rai-manual-v1.17.1_october_2019.pdf#page=669" class="req">Calculation
 * of PDPM Cognitive Level</a>
 */
public class CognitiveLevelLogic extends SnfDataVersionImpl<CognitiveLevel> {

  private final ClaimInfo claim;

  public CognitiveLevelLogic(ClaimInfo claim) {
    super(claim.getDataVersion());
    this.claim = claim;
  }

  /**
   * Determine the resident’s BIMS Summary Score on the MDS 3.0 based on the resident interview then
   * calculate the resident's PDPM cognitive level
   * <a href="doc-files/mds-3.0-rai-manual-v1.17.1_october_2019.pdf#page=669" class="req">step1</a>
   *
   * @return Cognitive level
   */
  protected CognitiveLevel step1(Assessment bimsSummaryScore) {
    CognitiveLevel result = null;
    if (hasSummaryScore(bimsSummaryScore)) {
      CognitiveLevelRow row = SnfTables.cognitiveLeveltable.values().stream()
          .flatMap(Collection::stream).filter(
              (item) -> item.isVersion(getDataVersion()) && item.isScore(
                  bimsSummaryScore.getValueInt())).findFirst().orElse(null);
      if (row != null) {
        result = SnfContext.trace(row.getCognitiveLevel());
      }
    }

    return result;
  }

  private static boolean hasSummaryScore(Assessment bimsSummaryScore) {
    boolean result = bimsSummaryScore != null;
    result = result && !bimsSummaryScore.isBlank();
    result = result && bimsSummaryScore.isAssessed();
    result = result && bimsSummaryScore.getValueInt() != Rai300.NONE_VALUE;
    return result;
  }

  /**
   * Determine the resident’s cognitive status based on the Staff Assessment for Mental Status for
   * the PDPM cognitive level <a href="doc-files/mds-3.0-rai-manual-v1.17.1_october_2019.pdf#page=669"
   * class= "req">step2To2a</a>
   *
   * @return Cognitive level
   */
  protected CognitiveLevel step2() {
    int c1000 = claim.getAssessmentValue(Rai300.C1000);

    CognitiveLevel result = step2a(c1000, claim::isComaAndNoActivities);

    // Step #2A
  if(result == null) {
      int b0700 = claim.getAssessmentValue(Rai300.B0700);
      int c0700 = claim.getAssessmentValue(Rai300.C0700);
      // Step #2B -> #2E
      result = step2bTo2e(b0700, c1000, c0700);
    }
    return SnfContext.trace(result);
  }

  protected static CognitiveLevel step2a(int c1000, Supplier<Boolean> isComa) {
    return c1000 == 3 || isComa.get() ? CognitiveLevel.SEVERELY_IMPAIRED : null;
  }

  /**
   * If the resident is not severely impaired based on Step A, then determine the resident’s Basic
   * Impairment Count and Severe Impairment Count <a href="doc-files/mds-3.0-rai-manual-v1.17.1_october_2019.pdf#page=670"
   * class= "req">step2bTo2e</a>
   *
   * @return Cognitive level
   */
  protected static CognitiveLevel step2bTo2e(int b0700, int c1000, int c0700) {
    CognitiveLevel result;

    // Step #2B
    int basicImpairmentCount = getBasicImpairmentCount(b0700, c1000, c0700);
    int severeImpairmentCount = getSevereImpairmentCount(b0700, c1000);

    // Step #2C
    if (severeImpairmentCount >= 1 && basicImpairmentCount >= 2) {
      result = CognitiveLevel.MODERATELY_IMPAIRED;
    } else
      // Step #2D
      if (basicImpairmentCount >= 1) {
        result = CognitiveLevel.MILDLY_IMPAIRED;
      } else {
        // Step #2E
        result = CognitiveLevel.COGNITIVELY_INTACT;
      }

    return SnfContext.trace(result);

  }

  private static int getSevereImpairmentCount(int b0700, int c1000) {
    int severeImpairmentCount = 0;
    severeImpairmentCount += (c1000 == 2) ? 1 : 0;
    severeImpairmentCount += (b0700 == 2 || b0700 == 3) ? 1 : 0;
    return severeImpairmentCount;
  }

  private static int getBasicImpairmentCount(int b0700, int c1000, int c0700) {
    int basicImpairmentCount = 0;
    basicImpairmentCount += (c1000 == 1 || c1000 == 2) ? 1 : 0;
    basicImpairmentCount += (b0700 >= 1 && b0700 <= 3) ? 1 : 0;
    basicImpairmentCount += (c0700 == 1) ? 1 : 0;
    return basicImpairmentCount;
  }

  @Override
  public CognitiveLevel exec() {
    CognitiveLevel result = step1(claim.getAssessment(Rai300.C0500));
    if (result == null) {
      result = step2();
    }

    return SnfContext.trace(result);

  }
}
