package main

import (
	"fmt"
)

func main() {
	part1, err := NewPart1([]int{
		5, 6, 6, 6, 6,
		3, 5, 6, 5, 6,
		1, 4, 6, -1, 6,
		3, 4, 2, 6, 5,
	})
	if err != nil {
		panic(err)
	}
	part2, err := NewPart2([]int{
		0, 6, 5, 1, 6,
		6, 1, 6, 5, 0,
		6, 4, 6, 3,
	})
	if err != nil {
		panic(err)
	}
	part3, err := NewPart3([]int{
		5, 0, 2, 1, 3,
		0, 1, 2, 5, 0,
		4, 2, 5, 0, 3,
		3, 5, 1,
	})
	if err != nil {
		panic(err)
	}

	for i, p := range []Part{part1, part2, part3} {
		p.Calculate()

		fmt.Printf("# Part %d\n\n", i+1)

		for k, v := range p.Scores() {
			fmt.Printf("%s: %f\n", k, v)
		}
		fmt.Printf("\n\n")
	}
}

type Spec map[string][]int

var Part1Spec = Spec{
	"Interference":      {2, 3, 4, 8, 9, 13, 14, 17, 19},
	"Support":           {5, 10, 15},
	"PainSeverity":      {1, 7, 12},
	"LifeControl":       {11, 16},
	"AffectiveDistress": {6, 18, 20},
}

var Part2Spec = Spec{
	"Negative":    {1, 4, 7, 10},
	"Solicitous":  {2, 5, 8, 11, 13, 14},
	"Distracting": {3, 6, 9, 12},
}

var Part3Spec = Spec{
	"HouseholdChores":        {1, 5, 9, 13, 17},
	"OutdoorWork":            {2, 6, 10, 14, 18},
	"ActivitiesAwayFromHome": {3, 7, 11, 15},
	"SocialActivities":       {4, 8, 12, 16},
	"GeneralActivity": {
		1, 2, 3, 4, 5,
		6, 7, 8, 9, 10,
		11, 12, 13, 14, 15,
		16, 17, 18,
	},
}

type Responses = []int
type Scores = map[string]float32

type Part interface {
	Responses() Responses
	Scores() Scores
	Calculate()
}

func NewPart1(responses []int) (Part, error) {
	if len(responses) != 20 {
		return nil, fmt.Errorf("Requires 20 responses")
	}
	p := &Part1{}
	p.PartBase.Responses = responses
	p.PartBase.Scores = make(Scores)

	return p, nil
}

type Part1 struct {
	PartBase
}

func (p *Part1) Responses() Responses {
	return p.PartBase.Responses
}

func (p *Part1) Scores() Scores {
	return p.PartBase.Scores
}

func (p *Part1) Calculate() {
	for k, v := range Part1Spec {
		if k == "AffectiveDistress" {
			continue
		}
		p.PartBase.Scores[k] = subMean(p.PartBase.Responses, v)
	}

	k := "AffectiveDistress"
	for _, q := range Part1Spec[k] {
		if q < 1 {
			panic(fmt.Sprintf("Cannot compute mean for nonpositive question %d", q))
		}
		if q == 6 {
			// Question 6 score is the complement of the response
			p.PartBase.Scores[k] +=
				6.0 - float32(p.Responses()[q-1])
			continue
		}
		p.PartBase.Scores[k] += float32(p.Responses()[q-1])
	}
	p.PartBase.Scores[k] /= float32(len(Part1Spec[k]))
}

func NewPart2(responses []int) (Part, error) {
	if len(responses) != 14 {
		return nil, fmt.Errorf("Requires 14 responses")
	}
	p := &Part2{}
	p.PartBase.Responses = responses
	p.PartBase.Scores = make(Scores)

	return p, nil
}

type Part2 struct {
	PartBase
}

func (p *Part2) Responses() Responses {
	return p.PartBase.Responses
}

func (p *Part2) Scores() Scores {
	return p.PartBase.Scores
}

func (p *Part2) Calculate() {
	p.PartBase.Calculate(Part2Spec)
}

func NewPart3(responses []int) (Part, error) {
	if len(responses) != 18 {
		return nil, fmt.Errorf("Requires 18 responses")
	}
	p := &Part3{}
	p.PartBase.Responses = responses
	p.PartBase.Scores = make(Scores)

	return p, nil
}

type Part3 struct {
	PartBase
}

func (p *Part3) Responses() Responses {
	return p.PartBase.Responses
}

func (p *Part3) Scores() Scores {
	return p.PartBase.Scores
}

func (p *Part3) Calculate() {
	p.PartBase.Calculate(Part3Spec)
}

type PartBase struct {
	Responses []int
	Scores    map[string]float32
}

func (p *PartBase) Calculate(spec Spec) {
	for k, v := range spec {
		p.Scores[k] = subMean(p.Responses, v)
	}
}

func subMean(s []int, qs []int) (mean float32) {
	for _, q := range qs {
		if q < 1 {
			panic(fmt.Sprintf("Cannot compute mean for nonpositive question %d", q))
		}
		mean += float32(s[q-1])
	}
	mean /= float32(len(qs))
	return
}