package trainedPredictionModels_test

import "seekia/internal/genetics/trainedPredictionModels"

import "testing"

import "seekia/internal/genetics/geneticReferences/polygenicDiseases"
import "seekia/internal/genetics/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.Fatal("InitializePolygenicDiseaseVariables failed: " + err.Error())
	}

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

	err = trainedPredictionModels.InitializeTrainedPredictionModels()
	if (err != nil){
		t.Fatal("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.Fatal("Prediction model not found: " + traitName)
			}

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

			traitLociList := traitObject.LociList

			computePrediction := func()error{

				// We have to try no more than 500 times, because sometimes we will create a fakeGenomeMap that does not
				//		possess any alleles which are present in the training data, thus, the prediction will fail

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

					testGenomeMap, err := getFakeGenomeMap(traitLociList)
					if (err != nil){
						return errors.New("getFakeGenomeMap failed: " + err.Error())
					}

					neuralNetworkExists, predictionIsPossible, _, _, _, _, err := trainedPredictionModels.GetNeuralNetworkDiscreteTraitPredictionFromGenomeMap(traitName, traitLociList, testGenomeMap)
					if (err != nil){
						return errors.New("GetNeuralNetworkDiscreteTraitPredictionFromGenomeMap failed: " + err.Error())
					}
					if (neuralNetworkExists == false){
						return errors.New("GetNeuralNetworkDiscreteTraitPredictionFromGenomeMap claims that neural network doesn't exist for trait: " + traitName)
					}
					if (predictionIsPossible == true){
						// Prediction was successful.
						return nil
					}
					// We try again
				}

				return errors.New("GetNeuralNetworkDiscreteTraitPredictionFromGenomeMap claims that prediction isn't possible.")
			}

			err = computePrediction()
			if (err != nil){
				t.Fatal(err)
			}
		}

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

		for _, attributeName := range numericAttributeNamesList{

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

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

				switch attributeName{

					case "Homosexualness",
						"Height":{

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

						traitLociList := traitObject.LociList

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

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

						diseaseLociList := diseaseObject.LociList

						return diseaseLociList, nil
					}
				}

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

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

			computePrediction := func()error{

				// We have to try no more than 500 times, because sometimes we will create a fakeGenomeMap that does not
				//		possess any alleles which are present in the training data, thus, the prediction will fail
				for i:=0; i < 500; i++{

					testGenomeMap, err := getFakeGenomeMap(attributeLociList)
					if (err != nil){
						return errors.New("getFakeGenomeMap failed: " + err.Error())
					}

					neuralNetworkExists, predictionIsPossible, _, _, _, _, err := trainedPredictionModels.GetNeuralNetworkNumericAttributePredictionFromGenomeMap(attributeName, attributeLociList, testGenomeMap)
					if (err != nil){
						return errors.New("GetNeuralNetworkNumericAttributePredictionFromGenomeMap failed: " + err.Error())
					}
					if (neuralNetworkExists == false){
						return errors.New("GetNeuralNetworkNumericAttributePredictionFromGenomeMap claims that neural network doesn't exist for attribute: " + attributeName)
					}
					if (predictionIsPossible == true){
						// Prediction was successful
						return nil
					}

					// Prediction failed, we try again.
				}

				return errors.New("GetNeuralNetworkNumericAttributePredictionFromGenomeMap failed to make a prediction.")
			}

			err = computePrediction()
			if (err != nil){
				t.Fatal(err)
			}
		}
	}
}


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
}

