package no.valg.valgnatt.backend.mandatberegning.prognose;

import java.util.Map;

import no.valg.valgnatt.backend.mandatberegning.domene.KommuneKretsParti;
import no.valg.valgnatt.backend.mandatberegning.domene.Kvotient;
import no.valg.valgnatt.common.ValgnattException;
import no.valg.valgnatt.database.Konstanter;

import org.apache.logging.log4j.LogManager;
import org.apache.logging.log4j.Logger;
import org.springframework.stereotype.Component;

import cern.colt.matrix.DoubleMatrix2D;
import cern.colt.matrix.impl.DenseDoubleMatrix2D;
import cern.colt.matrix.linalg.Algebra;
import cern.colt.matrix.linalg.QRDecomposition;

/**
 * MSI mai 2016:
 * Denne klassen har som formål å beregne de såkalte betakoeffisientene/verdier (BK) som brukes
 * i prognoseberegningen. BK kunne i prinsippet vært satt til konstante verdier, og man kunne sluppet
 * hele beregningen i denne klassen, men noen simuleringer og tester viser at med
 * beregnede BK så kan man få mer presise prognoser.
 *
 * Hvordan brukes så BK'ene? Regnestykke(ne) ser slik ut:
 *
 * sampleAndel(KRF) = BK1*prediktivAndel1(KRF) + BK2*prediktivAndel2(KRF) + ... (dersom pred.-array er større)
 * sampleAndel(SV) = BK1*prediktivAndel1(SV) + BK2*prediktivAndel2(SV) + ... (dersom pred.-array er større)
 * sampleAndel(SP) = BK1*prediktivAndel1(SP) + BK2*prediktivAndel2(SP) + ... (dersom pred.-array er større)
 * ....
 * ....
 * (En linje pr parti/krets)
 *
 * Poenget er at matrisekalkulasjonen (ser du at regnestykkene over er en matrise?) forsøker å finne
 * de BK'ene som gir den beste sammenheng/løsning på regnestykkene over. Legg merke til at det er de samme
 * BK'ene som går igjen i hver linje. Man kan altså tenke seg følgende eks (anta at linjene tilsvarer
 * partiene vi brukte over):
 *
 * 0.045 = BK1*0.05 + BK2*0.06
 * 0.040 = BK1*0.03 + BK2*0.05
 * 0.074 = BK1*0.04 + BK2*0.05
 *
 * Jeg vet ikke eksakt hva slags BK-verdier vi da får, men la oss for enkelhets skyld si at BK1 = 0.345 og
 * BK2 = 0.471. I så tilfelle kan man raskt slå fast at regnestykkene ikke "går opp". Dette er da heller ikke
 * poenget. BK'ene angir nå de koeffisientene som gjør at ALLE disse regnestykkene totalt sett kommer nærmest
 * sluttresultatet. Dermed angir BK *sammenhengen* mellom prediktive-andeler (= FHS og historiske valg) og
 * sluttresultatet (Sample). Denne sammenhengen gjør det mulig å gå fra FHS+historiske valg til å finne
 * sluttrestaultat for hele landet (eller hele fylket/området).
 */
@Component
public final strictfp class MatriseKalkulasjon { // strictfp vil sikrer at system/enhetstester fungerer på alle plattformer
	private static final Logger LOG = LogManager.getLogger(MatriseKalkulasjon.class);
	private static final Algebra algebra = new Algebra(Double.MIN_VALUE);
	private static double apriroBkFhs = Konstanter.DEFAULT_APRIORI_BK_FHS;
	private static double aprioriBkHs1 = Konstanter.DEFAULT_APRIORI_BK_HS1;
	private static double aprioriBkHs2 = Konstanter.DEFAULT_APRIORI_BK_HS2;

	private MatriseKalkulasjon() {
	}


	@SafeVarargs
	static double[] beregnBetaverdier(Map<KommuneKretsParti, Kvotient> sampleAndeler, Map<KommuneKretsParti, Kvotient>... prediktiveAndelerArr) {

		// Apriori-betaverdier er utgangs/default-verdier for beta-koeffisientene. Dvs. at i begynnelsen av opptellingen
		// med få stemmer inne, så vil betaverdiene bli "dratt" mot apriori-verdiene. En måte å tenke på dette er at vi
		// mener at det kan være fornuftig å vekte FHS til 0.5 i utgangspunktet, og på den måten tillegge FHS-komponenten
		// stor verdi til å begynne med. Skulle dette vise seg å være feil, så vil denne beta-koeffisienten for FHS
		// bli modifisert utover i valget.
		double[] aprioriBetaverdier = getAprioriBetaVerdier(prediktiveAndelerArr.length);
		if (Prognose.STATISK_BETA_KOEFFISIENTER) {
			return getAprioriBetaverdier(aprioriBetaverdier);
		}

		try {
			// Z-matrisen har ferdigOpptaltAnder.size() rader og (prediktiveAndelerPrKKPArr.length + 1) kolonner
			DoubleMatrix2D zMatrix = lagZMatrise(sampleAndeler, prediktiveAndelerArr);
			// Akkurat hva en QR-dekomposisjon er, er uvisst. Men dersom det mangler data
			// (f.eks. at en kolonne har bare 0-verdier), så kan man få en manglende
			// "Full Rank" - dvs. at man mangler data for å gjennomføre matrisekalkulasjonen.
			QRDecomposition qrDecomposition = beregnQRDecomposition(zMatrix);
			// U-matrisen består simpelthen av alle andelene fra ferdigOpptaltAndeler, altså 1 kolonne bred
			DoubleMatrix2D uMatrise = lagUMatrise(sampleAndeler);
			// Så fortsetter kalkulasjonen, i henhold til prognosedokumentet (se nederst på side 2 og øverst på side 3).
			DoubleMatrix2D bMatrise = qrDecomposition.solve(uMatrise);
			DoubleMatrix2D zbMatrise = algebra.mult(zMatrix, bMatrise);
			// Sigmakvadratet er kun viktig for å beregne standardavviket senere
			Double sigmaKvadrat = beregnSigmaKvadrat(uMatrise, zbMatrise, sampleAndeler.size(), prediktiveAndelerArr.length + 1);
			// På en eller annen måte har nå betaverdiene blitt beregnet, men disse må nå
			// "trekkes" mot apriori-verdiene i henhold til formelen gitt litt nede på side 3 i prognosedokumentet
			DoubleMatrix2D betaVerdier = beregnBetaVerdier(bMatrise, sigmaKvadrat, qrDecomposition, lagAprioriMatrise(aprioriBetaverdier));
			// javel, så har vi beta-verdiene da - bare litt konvertering av datastruktur, så er vi i mål
			double[] betas = new double[betaVerdier.rows()];
			for (int i = 0; i < betas.length; i++) {
				betas[i] = betaVerdier.get(i, 0);
			}
			return betas;
		} catch (Exception e) {
			LOG.warn("Må bruke apriori beta-verdier, fordi vi har for få kretser til å gjøre matrisekalkulasjonen: " + e);
			// Dersom vi har for lite data til å løse matriseligningen, så bruker vi simpelthen apriori-verdiene. Ettersom flere
			// stemmer blir talt opp, så vil man mer sannsynlig få nok data til å løse matriselignen, og dermed få bedre beta-
			// verdier.
			return getAprioriBetaverdier(aprioriBetaverdier);
		}
	}

	private static double[] getAprioriBetaverdier(double[] aprioriBetaverdier) {
		double[] betas = new double[aprioriBetaverdier.length + 1];
		betas[0] = 0;
		System.arraycopy(aprioriBetaverdier, 0, betas, 1, betas.length - 1);
		return betas;
	}

	private static double[] getAprioriBetaVerdier(int antallPrediktiveStemmetall) {
		double[] aprioriBetaverdier;
		if (antallPrediktiveStemmetall == 2) { // Bare siste/nestsiste-valg
			aprioriBetaverdier = new double[]{aprioriBkHs1, aprioriBkHs2};
		} else { // FHS + siste/nestsiste-valg
			aprioriBetaverdier = new double[]{apriroBkFhs, aprioriBkHs1, aprioriBkHs2};
		}
		return aprioriBetaverdier;
	}

	private static DoubleMatrix2D lagAprioriMatrise(final double[] aprioriValues) {
		final DoubleMatrix2D result = new DenseDoubleMatrix2D(aprioriValues.length + 1, 1);
		result.set(0, 0, 0);
		for (int i = 1; i <= aprioriValues.length; i++) {
			result.set(i, 0, aprioriValues[i - 1]);
		}
		return result;
	}

	/**
	 * MSI jun 2016:
	 * Akkurat hva som foregår i denne metoden kan jeg ikke gi noen god forklaring på. Her er det en del
	 * matrisematematikk som beregner beta-verdiene.
	 */
	@java.lang.SuppressWarnings("squid:S1244")
	private static DoubleMatrix2D beregnBetaVerdier(DoubleMatrix2D bMatrise, Double sigmaKvadrat, QRDecomposition qrDecomposition, DoubleMatrix2D aprioriBetas) {
		DoubleMatrix2D kMatrix = algebra.mult(algebra.transpose(qrDecomposition.getR()), qrDecomposition.getR());
		if (algebra.det(kMatrix) == 0) {
			LOG.debug("Determinant of k Matrix is 0, using apriori values instead!");
			return aprioriBetas;
		}
		DoubleMatrix2D kInverse = algebra.inverse(kMatrix);
		DoubleMatrix2D tempBetas = new DenseDoubleMatrix2D(kMatrix.rows(), 1);
		Double factor = sigmaKvadrat / (Konstanter.TAU * Konstanter.TAU);
		DoubleMatrix2D kBetaMatrix = new DenseDoubleMatrix2D(kMatrix.rows(), kMatrix.columns());
		for (int row = 0; row < kMatrix.rows(); row++) {
			tempBetas.set(row, 0, bMatrise.get(row, 0) - aprioriBetas.get(row, 0));
			for (int col = 0; col < kMatrix.columns(); col++) {
				double val = (col == row) ? 1.0 : 0.0;
				val += (factor * kInverse.get(row, col));
				kBetaMatrix.set(row, col, val);
			}
		}
		DoubleMatrix2D deltaBetas = algebra.solve(kBetaMatrix, tempBetas);
		return summerMatriser(deltaBetas, aprioriBetas);
	}

	private static DoubleMatrix2D summerMatriser(DoubleMatrix2D deltaBetas, DoubleMatrix2D aprioriBetas) {
		DoubleMatrix2D result = new DenseDoubleMatrix2D(deltaBetas.rows(), deltaBetas.columns());
		for (int row = 0; row < deltaBetas.rows(); row++) {
			for (int col = 0; col < deltaBetas.columns(); col++) {
				result.set(row, col, aprioriBetas.get(row, col) + deltaBetas.get(row, col));
			}
		}
		return result;
	}

	private static Double beregnSigmaKvadrat(DoubleMatrix2D uMatrise, DoubleMatrix2D zbMatrise, int n, int p) {
		double temp = 0.0;
		for (int row = 0; row < uMatrise.rows(); row++) {
			for (int column = 0; column < uMatrise.columns(); column++) {
				temp = temp + Math.pow(uMatrise.get(row, column) - zbMatrise.get(row, column), 2);
			}
		}
		if (n - p == 0) {
			throw new ValgnattException("For få kretser er klare, slik at man kan utføre matriseberegningen - bruker default-verdier");
		}
		return temp / (double) (n - p);
	}

	private static QRDecomposition beregnQRDecomposition(DoubleMatrix2D matrise) {
		QRDecomposition result = new QRDecomposition(matrise);
		if (result.hasFullRank()) {
			return result;
		} else {
			throw new ValgnattException("Mangler såkalt 'Full Rank' for følgende matrise. Kan være at en kolonne inneholder bare 0'ere -> feil i stemmetall?\n" +
					matrise.toString());
		}
	}

	@SafeVarargs
	private static DoubleMatrix2D lagZMatrise(Map<KommuneKretsParti, Kvotient> ferdigOpptaltAndeler, Map<KommuneKretsParti, Kvotient>... prediktiveAndeler) {
		DoubleMatrix2D result = new DenseDoubleMatrix2D(ferdigOpptaltAndeler.size(), prediktiveAndeler.length + 1);
		int row = 0;
		for (KommuneKretsParti mdp : ferdigOpptaltAndeler.keySet()) {
			result.setQuick(row, 0, 1); // Ikke spør meg, men første kolonne skal være 1
			int columnNumber = 1;
			for (Map<KommuneKretsParti, Kvotient> prediktivAndelerMap : prediktiveAndeler) {
				result.setQuick(row, columnNumber++, prediktivAndelerMap.get(mdp) == null ? 0 : prediktivAndelerMap.get(mdp).getKvotient().doubleValue());
			}
			row++;
		}
		return result;
	}

	private static DoubleMatrix2D lagUMatrise(Map<KommuneKretsParti, Kvotient> ferdigOpptaltAndeler) {
		DoubleMatrix2D result = new DenseDoubleMatrix2D(ferdigOpptaltAndeler.size(), 1);
		int row = 0;
		for (KommuneKretsParti mdp : ferdigOpptaltAndeler.keySet()) {
			result.setQuick(row, 0, ferdigOpptaltAndeler.get(mdp).getKvotient().doubleValue());
			row++;
		}
		return result;
	}

	static void setAprioriBetaVerdier(double aprioriFhs, double aprioriHs1, double aprioriHs2) {
		apriroBkFhs = aprioriFhs;
		aprioriBkHs1 = aprioriHs1;
		aprioriBkHs2 = aprioriHs2;
	}
}
