package trainedPredictionModels_test

import "seekia/resources/trainedPredictionModels"

import "testing"

import "seekia/resources/geneticReferences/polygenicDiseases"
import "seekia/resources/geneticReferences/traits"

import "seekia/internal/genetics/locusValue"
import "seekia/internal/helpers"

import "errors"


func TestTrainedPredictionModels(t *testing.T){

	err := polygenicDiseases.InitializePolygenicDiseaseVariables()
	if (err != nil) {
		t.Fatalf("InitializePolygenicDiseaseVariables failed: " + err.Error())
	}

	err = traits.InitializeTraitVariables()
	if (err != nil) {
		t.Fatalf("InitializeTraitVariables failed: " + err.Error())
	}

	err = trainedPredictionModels.InitializeTrainedPredictionModels()
	if (err != nil){
		t.Fatalf("InitializeTrainedPredictionModels failed: " + err.Error())
	}

	for i:=0; i < 10; i++{

		discreteTraitNamesList := []string{"Eye Color", "Lactose Tolerance"}

		for _, traitName := range discreteTraitNamesList{

			modelExists := trainedPredictionModels.CheckIfAttributeNeuralNetworkExists(traitName)
			if (modelExists == false){
				t.Fatalf("Prediction model not found: " + traitName)
			}

			traitObject, err := traits.GetTraitObject(traitName)
			if (err != nil) {
				t.Fatalf("GetTraitObject failed: " + err.Error())
			}

			traitLociList := traitObject.LociList

			testGenomeMap, err := getFakeGenomeMap(traitLociList)
			if (err != nil){
				t.Fatalf("getFakeGenomeMap failed: " + err.Error())
			}

			neuralNetworkExists, predictionIsPossible, _, _, _, _, err := trainedPredictionModels.GetNeuralNetworkDiscreteTraitPredictionFromGenomeMap(traitName, traitLociList, testGenomeMap)
			if (err != nil){
				t.Fatalf("GetNeuralNetworkDiscreteTraitPredictionFromGenomeMap failed: " + err.Error())
			}
			if (neuralNetworkExists == false){
				t.Fatalf("GetNeuralNetworkDiscreteTraitPredictionFromGenomeMap claims that neural network doesn't exist for trait: " + traitName)
			}
			if (predictionIsPossible == false){
				t.Fatalf("GetNeuralNetworkDiscreteTraitPredictionFromGenomeMap claims that prediction isn't possible.")
			}
		}

		numericAttributeNamesList := []string{"Height", "Autism", "Obesity", "Homosexualness"}

		for _, attributeName := range numericAttributeNamesList{

			modelExists := trainedPredictionModels.CheckIfAttributeNeuralNetworkExists(attributeName)
			if (modelExists == false){
				t.Fatalf("Prediction model not found: " + attributeName)
			}

			getAttributeLociList := func()([]int64, error){

				switch attributeName{

					case "Homosexualness",
						"Height":{

						traitObject, err := traits.GetTraitObject(attributeName)
						if (err != nil) {
							t.Fatalf("GetTraitObject failed: " + attributeName)
						}

						traitLociList := traitObject.LociList

						return traitLociList, nil
					}
					case "Obesity",
						"Autism":{

						diseaseObject, err := polygenicDiseases.GetPolygenicDiseaseObject(attributeName)
						if (err != nil){
							t.Fatalf("GetPolygenicDiesaseObject failed: " + err.Error())
						}

						diseaseLociList := diseaseObject.LociList

						return diseaseLociList, nil
					}
				}

				return nil, errors.New("Unknown attributeName: " + attributeName)
			}

			attributeLociList, err := getAttributeLociList()
			if (err != nil){
				t.Fatalf(err.Error())
			}

			testGenomeMap, err := getFakeGenomeMap(attributeLociList)
			if (err != nil){
				t.Fatalf("getFakeGenomeMap failed: " + err.Error())
			}

			neuralNetworkExists, predictionIsPossible, _, _, _, _, err := trainedPredictionModels.GetNeuralNetworkNumericAttributePredictionFromGenomeMap(attributeName, attributeLociList, testGenomeMap)
			if (err != nil){
				t.Fatalf("GetNeuralNetworkNumericAttributePredictionFromGenomeMap failed: " + err.Error())
			}
			if (neuralNetworkExists == false){
				t.Fatalf("GetNeuralNetworkNumericAttributePredictionFromGenomeMap claims that neural network doesn't exist for attribute: " + attributeName)
			}
			if (predictionIsPossible == false){
				t.Fatalf("GetNeuralNetworkNumericAttributePredictionFromGenomeMap claims that prediction isn't possible.")
			}
		}
	}
}


func getFakeGenomeMap(lociList []int64)(map[int64]locusValue.LocusValue, error){

	// We create a fake genome map

	testGenomeMap := make(map[int64]locusValue.LocusValue)

	for index, rsID := range lociList{

		if (index != 0){
			// We always include the first locus

			// We will include approximately 80% of the locations in the genome
			includeLocusBool, err := helpers.GetRandomBoolWithProbability(0.8)
			if (err != nil){ return nil, err }
			if (includeLocusBool == false){
				continue
			}
		}

		locusIsPhasedBool := helpers.GetRandomBool()

		randomAllele1, err := helpers.GetRandomItemFromList([]string{"C", "A", "T", "G", "I", "D"})
		if (err != nil){ return nil, err }
		randomAllele2, err := helpers.GetRandomItemFromList([]string{"C", "A", "T", "G", "I", "D"})			
		if (err != nil){ return nil, err }

		newLocusValue := locusValue.LocusValue{
			Base1Value: randomAllele1,
			Base2Value: randomAllele2,
			LocusIsPhased: locusIsPhasedBool,
		}

		testGenomeMap[rsID] = newLocusValue
	}

	return testGenomeMap, nil
}