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.Pair;
import gov.cms.grouper.snf.lego.SnfCache;
import gov.cms.grouper.snf.lego.SnfComparator;
import gov.cms.grouper.snf.lego.SnfUtils;
import gov.cms.grouper.snf.lego.SnfDataVersionImpl;
import gov.cms.grouper.snf.model.Assessment;
import gov.cms.grouper.snf.model.CognitiveLevel;
import gov.cms.grouper.snf.model.reader.Rai300;
import gov.cms.grouper.snf.model.table.CognitiveLevelRow;
import gov.cms.grouper.snf.util.ClaimInfo;
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;
  private final Supplier<Integer> c1000Supplier;
  private final Supplier<Integer> b0700Supplier;
  private final Supplier<Integer> c0700Supplier;
  private final Supplier<Integer> b0100Supplier;
  private final Supplier<Assessment> bimsSummaryScore;

  public CognitiveLevelLogic(ClaimInfo claim) {
    this(claim, null, null, null, null);
  }

  protected CognitiveLevelLogic(ClaimInfo claim, Supplier<Integer> c1000Supplier,
      Supplier<Integer> b0700Supplier, Supplier<Integer> c0700Supplier,
      Supplier<Integer> b0100Supplier) {
    super(claim.getDataVersion());
    this.claim = claim;
    this.c1000Supplier = SnfUtils.nullCheck(c1000Supplier, this.of(Rai300.C1000), c1000Supplier);
    this.b0700Supplier = SnfUtils.nullCheck(b0700Supplier, this.of(Rai300.B0700), b0700Supplier);
    this.c0700Supplier = SnfUtils.nullCheck(c0700Supplier, this.of(Rai300.C0700), c0700Supplier);
    this.b0100Supplier = SnfUtils.nullCheck(b0100Supplier, this.of(Rai300.B0100), b0100Supplier);

    this.bimsSummaryScore = this.assessmentOf(Rai300.C0500);
  }

  /**
   * 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(Supplier<CognitiveLevel> next,
      Supplier<Assessment> bimsSummaryScoreSupplier) {
    CognitiveLevel result = null;
    final Assessment bimsSummaryAssessment = bimsSummaryScoreSupplier.get();
    if (bimsSummaryAssessment != null && !bimsSummaryAssessment.isBlank()
        && bimsSummaryAssessment.isAssessed() && bimsSummaryAssessment.getValueInt() != 99) {
      CognitiveLevelRow row =
          SnfTables.get(SnfTables.cognitiveLeveltable, CognitiveLevelRow.selector,
              Pair.of(this.getDataVersion(), bimsSummaryAssessment.getValueInt()));
      if (row != null) {
        result = row.getCognitiveLevel();
      }
    }

    if (result == null) {
      result = next.get();
    }

    return SnfContext.trace(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 step2To2a(Supplier<CognitiveLevel> next, Supplier<Integer> c1000Supplier,
      Supplier<Boolean> isComaAndNoActivitiesSupplier) {
    CognitiveLevel result = null;

    // Step #2A
    if (c1000Supplier.get() == 3 || isComaAndNoActivitiesSupplier.get()) {
      result = CognitiveLevel.SEVERELY_IMPARED;
    }

    if (result == null) {
      result = next.get();
    }

    return SnfContext.trace(result);

  }

  /**
   * 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 CognitiveLevel step2bTo2e(Supplier<Integer> b0700Supplier,
      Supplier<Integer> c1000Supplier, Supplier<Integer> c0700Supplier) {
    CognitiveLevel result = null;

    final int b0700 = b0700Supplier.get();
    final int c1000 = c1000Supplier.get();
    final int c0700 = c0700Supplier.get();
    // Step #2B
    int basicImpairmentCount = 0;
    basicImpairmentCount += (c1000 == 1 || c1000 == 2) ? 1 : 0;
    basicImpairmentCount += (b0700 >= 1 && b0700 <= 3) ? 1 : 0;
    basicImpairmentCount += (c0700 == 1) ? 1 : 0;

    int severeImpairmentCount = 0;
    severeImpairmentCount += (c1000 == 2) ? 1 : 0;
    severeImpairmentCount += (b0700 == 2 || b0700 == 3) ? 1 : 0;

    if (SnfComparator.betweenInclusive(1, severeImpairmentCount, 2)
        && SnfComparator.betweenInclusive(2, basicImpairmentCount, 3)) {
      // Step #2C
      result = CognitiveLevel.MODERATELY_IMPAIRED;
    } else if ((basicImpairmentCount == 1
        && SnfComparator.betweenInclusive(0, severeImpairmentCount, 2))
        || (SnfComparator.betweenInclusive(2, basicImpairmentCount, 3)
            && severeImpairmentCount == 0)) {
      // Step #2D
      result = CognitiveLevel.MILDLY_IMPAIRED;
    } else if (basicImpairmentCount == 0 && severeImpairmentCount == 0) {
      // Step #2E
      result = CognitiveLevel.COGNITIVELY_INTACT;
    }

    return SnfContext.trace(result);

  }

  @Override
  public CognitiveLevel exec() {
    Supplier<CognitiveLevel> step2bTo2e =
        () -> this.step2bTo2e(this.b0700Supplier, this.c1000Supplier, this.c0700Supplier);
    Supplier<CognitiveLevel> step2To2a = () -> this.step2To2a(step2bTo2e, this.c1000Supplier,
        () -> this.claim.isComaAndNoActivities(this.b0100Supplier));
    CognitiveLevel result = this.step1(step2To2a, this.bimsSummaryScore);

    return SnfContext.trace(result);

  }

  public Supplier<Integer> of(Rai300 input) {
    return SnfCache.of(() -> this.claim.getAssessmentValue(input));
  }

  public Supplier<Assessment> assessmentOf(Rai300 input) {
    return () -> this.claim.getAssessment(input);
  }

}
