package main

import (
	"errors"
	"math"
	"strconv"

	"gonum.org/v1/gonum/graph"
	"gonum.org/v1/gonum/graph/encoding"
	"gonum.org/v1/gonum/graph/simple"
	"gonum.org/v1/gonum/graph/topo"
)

type projectNetwork struct {
	*simple.DirectedGraph
}

func newProjectNetwork() *projectNetwork {
	return &projectNetwork{DirectedGraph: simple.NewDirectedGraph()}
}

func (g *projectNetwork) NewNode() graph.Node {
	return &activity{Node: g.DirectedGraph.NewNode()}
}

func (g *projectNetwork) calculateTimes() error {
	sorted, err := topo.Sort(g)
	if err != nil {
		return err
	}

	toposort := make([]*activity, len(sorted))
	for i := range sorted {
		toposort[i] = sorted[i].(*activity)
	}

	g.forwardPass(toposort)
	g.backwardPass(toposort)

	return nil
}

func (g *projectNetwork) forwardPass(toposort []*activity) {
	for _, a := range toposort {
		g.setEarlyTimes(a)
	}
}

func (g *projectNetwork) backwardPass(toposort []*activity) {
	n := len(toposort)

	end := toposort[n-1]
	end.lateFinish = end.earlyFinish
	end.lateStart = end.lateFinish - end.duration

	for i := n - 2; i >= 0; i-- {
		a := toposort[i]
		g.setLateTimes(a)
	}
}

func (g *projectNetwork) setEarlyTimes(a *activity) {
	es := 0.0
	predecessors := g.To(a.ID())
	for predecessors.Next() {
		p := predecessors.Node().(*activity)
		es = math.Max(es, p.earlyFinish)
	}
	a.earlyStart = es
	a.earlyFinish = a.earlyStart + a.duration
}

func (g *projectNetwork) setLateTimes(a *activity) {
	lf := math.Inf(1)
	successors := g.From(a.ID())
	for successors.Next() {
		s := successors.Node().(*activity)
		lf = math.Min(lf, s.lateStart)
	}
	a.lateFinish = lf
	a.lateStart = a.lateFinish - a.duration
	a.slack = a.lateFinish - a.earlyFinish
}

type activity struct {
	graph.Node
	dotID       string
	duration    float64
	earlyStart  float64
	earlyFinish float64
	lateStart   float64
	lateFinish  float64
	slack       float64
}

func (a *activity) Attributes() []encoding.Attribute {
	return []encoding.Attribute{
		{"DUR", strconv.FormatFloat(a.duration, 'f', 2, 64)},
		{"ES", strconv.FormatFloat(a.earlyStart, 'f', 2, 64)},
		{"EF", strconv.FormatFloat(a.earlyFinish, 'f', 2, 64)},
		{"LS", strconv.FormatFloat(a.lateStart, 'f', 2, 64)},
		{"LF", strconv.FormatFloat(a.lateFinish, 'f', 2, 64)},
		{"SLACK", strconv.FormatFloat(a.slack, 'f', 2, 64)},
	}
}

func (a *activity) SetAttribute(attr encoding.Attribute) error {
	var err error
	switch attr.Key {
	case "DUR":
		a.duration, err = strconv.ParseFloat(attr.Value, 64)
	default:
		err = errors.New("can't set attribute " + attr.Key)
	}
	return err
}

func (a *activity) DOTID() string {
	return a.dotID
}

func (a *activity) SetDOTID(id string) {
	a.dotID = id
}