package transformers

import (
	"fmt"
	"os"
	"strings"

	"github.com/johnkerl/miller/v6/pkg/cli"
	"github.com/johnkerl/miller/v6/pkg/lib"
	"github.com/johnkerl/miller/v6/pkg/mlrval"
	"github.com/johnkerl/miller/v6/pkg/transformers/utils"
	"github.com/johnkerl/miller/v6/pkg/types"
)

const verbNameStats2 = "stats2"

// For joining "x" and "y" into "x...y" for map keys. "," is another natural choice but would break
// if we were ever asked to process field names with commas in them.
const stats2KeySeparator = "\001"

var Stats2Setup = TransformerSetup{
	Verb:         verbNameStats2,
	UsageFunc:    transformerStats2Usage,
	ParseCLIFunc: transformerStats2ParseCLI,
	IgnoresInput: false,
}

func transformerStats2Usage(
	o *os.File,
) {
	argv0 := "mlr"
	verb := verbNameStats2

	fmt.Fprintf(o, "Usage: %s %s [options]\n", argv0, verb)
	fmt.Fprintf(o, "Computes bivariate statistics for one or more given field-name pairs,\n")
	fmt.Fprintf(o, "accumulated across the input record stream.\n")
	fmt.Fprintf(o, "-a {linreg-ols,corr,...}  Names of accumulators: one or more of:\n")

	utils.ListStats2Accumulators(o)

	fmt.Fprintf(o, "-f {a,b,c,d}   Value-field name-pairs on which to compute statistics.\n")
	fmt.Fprintf(o, "               There must be an even number of names.\n")
	fmt.Fprintf(o, "-g {e,f,g}     Optional group-by-field names.\n")
	fmt.Fprintf(o, "-v             Print additional output for linreg-pca.\n")
	fmt.Fprintf(o, "-s             Print iterative stats. Useful in tail -f contexts, in which\n")
	fmt.Fprintf(o, "               case please avoid pprint-format output since end of input\n")
	fmt.Fprintf(o, "               stream will never be seen. Likewise, if input is coming from\n")
	fmt.Fprintf(o, "               `tail -f`, be sure to use `--records-per-batch 1`.\n")
	fmt.Fprintf(o, "--fit          Rather than printing regression parameters, applies them to\n")
	fmt.Fprintf(o, "               the input data to compute new fit fields. All input records are\n")
	fmt.Fprintf(o, "               held in memory until end of input stream. Has effect only for\n")
	fmt.Fprintf(o, "               linreg-ols, linreg-pca, and logireg.\n")
	fmt.Fprintf(o, "Only one of -s or --fit may be used.\n")
	fmt.Fprintf(o, "Example: %s %s -a linreg-pca -f x,y\n", argv0, verb)
	fmt.Fprintf(o, "Example: %s %s -a linreg-ols,r2 -f x,y -g size,shape\n", argv0, verb)
	fmt.Fprintf(o, "Example: %s %s -a corr -f x,y\n", argv0, verb)
}

func transformerStats2ParseCLI(
	pargi *int,
	argc int,
	args []string,
	_ *cli.TOptions,
	doConstruct bool, // false for first pass of CLI-parse, true for second pass
) (RecordTransformer, error) {

	// Skip the verb name from the current spot in the mlr command line
	argi := *pargi
	verb := args[argi]
	argi++

	var accumulatorNameList []string = nil
	var valueFieldNameList []string = nil
	groupByFieldNameList := []string{}
	doVerbose := false
	doIterativeStats := false
	doHoldAndFit := false

	var err error
	for argi < argc /* variable increment: 1 or 2 depending on flag */ {
		opt := args[argi]
		if !strings.HasPrefix(opt, "-") {
			break // No more flag options to process
		}
		if args[argi] == "--" {
			break // All transformers must do this so main-flags can follow verb-flags
		}
		argi++

		if opt == "-h" || opt == "--help" {
			transformerStats2Usage(os.Stdout)
			return nil, cli.ErrHelpRequested

		} else if opt == "-a" {
			accumulatorNameList, err = cli.VerbGetStringArrayArg(verb, opt, args, &argi, argc)
			if err != nil {
				return nil, err
			}

		} else if opt == "-f" {
			valueFieldNameList, err = cli.VerbGetStringArrayArg(verb, opt, args, &argi, argc)
			if err != nil {
				return nil, err
			}

		} else if opt == "-g" {
			groupByFieldNameList, err = cli.VerbGetStringArrayArg(verb, opt, args, &argi, argc)
			if err != nil {
				return nil, err
			}

		} else if opt == "-v" {
			doVerbose = true

		} else if opt == "-s" {
			doIterativeStats = true

		} else if opt == "--fit" {
			doHoldAndFit = true

		} else if opt == "-S" {
			// No-op pass-through for backward compatibility with Miller 5

		} else if opt == "-F" {
			// The -F flag isn't used for stats2: all arithmetic here is
			// floating-point. Yet it is supported for step and stats1 for all
			// applicable stats1/step accumulators, so we accept here as well
			// for all applicable stats2 accumulators (i.e. none of them).

		} else {
			return nil, cli.VerbErrorf(verb, "option \"%s\" not recognized", opt)
		}
	}

	if doIterativeStats && doHoldAndFit {
		return nil, cli.VerbErrorf(verb, "cannot combine -I and -H")
	}
	if accumulatorNameList == nil {
		return nil, cli.VerbErrorf(verb, "-a option is required")
	}
	if valueFieldNameList == nil {
		return nil, cli.VerbErrorf(verb, "-f option is required")
	}
	if len(valueFieldNameList)%2 != 0 {
		return nil, cli.VerbErrorf(verb, "argument to -f must have even number of fields")
	}

	*pargi = argi
	if !doConstruct { // All transformers must do this for main command-line parsing
		return nil, nil
	}

	transformer, err := NewTransformerStats2(
		accumulatorNameList,
		valueFieldNameList,
		groupByFieldNameList,
		doVerbose,
		doIterativeStats,
		doHoldAndFit,
	)
	if err != nil {
		return nil, err
	}

	return transformer, nil
}

type TransformerStats2 struct {
	// Input:
	accumulatorNameList  []string
	valueFieldNameList   []string
	groupByFieldNameList []string

	doVerbose        bool
	doIterativeStats bool
	doHoldAndFit     bool

	// State:
	accumulatorFactory *utils.Stats2AccumulatorFactory

	// Accumulators are indexed by
	//   groupByFieldName . value1FieldName+sep+value2FieldName . accumulatorName . accumulator object
	// This would be
	//   namedAccumulators map[string]map[string]map[string]IStats2Accumulator
	// except we need maps that preserve insertion order.
	namedAccumulators *lib.OrderedMap[*lib.OrderedMap[*lib.OrderedMap[utils.IStats2Accumulator]]]

	groupingKeysToGroupByFieldValues *lib.OrderedMap[[]*mlrval.Mlrval]

	// For hold-and-fit:
	// ordered map from grouping-key to list of RecordAndContext
	recordGroups *lib.OrderedMap[*[]*types.RecordAndContext]
}

func NewTransformerStats2(
	accumulatorNameList []string,
	valueFieldNameList []string,
	groupByFieldNameList []string,
	doVerbose bool,
	doIterativeStats bool,
	doHoldAndFit bool,
) (*TransformerStats2, error) {
	for _, name := range accumulatorNameList {
		if !utils.ValidateStats2AccumulatorName(name) {
			return nil, fmt.Errorf(`mlr stats2: accumulator "%s" not found`, name)
		}
	}

	tr := &TransformerStats2{
		accumulatorNameList:              accumulatorNameList,
		valueFieldNameList:               valueFieldNameList,
		groupByFieldNameList:             groupByFieldNameList,
		doVerbose:                        doVerbose,
		doIterativeStats:                 doIterativeStats,
		doHoldAndFit:                     doHoldAndFit,
		accumulatorFactory:               utils.NewStats2AccumulatorFactory(),
		namedAccumulators:                lib.NewOrderedMap[*lib.OrderedMap[*lib.OrderedMap[utils.IStats2Accumulator]]](),
		groupingKeysToGroupByFieldValues: lib.NewOrderedMap[[]*mlrval.Mlrval](),
		recordGroups:                     lib.NewOrderedMap[*[]*types.RecordAndContext](),
	}
	return tr, nil
}

// Given: accumulate corr,cov on values x,y group by a,b.
// Example input:       Example output:
//   a b x y            a b x_corr x_cov y_corr y_cov
//   s t 1 2            s t 2       6    2      8
//   u v 3 4            u v 1       3    1      4
//   s t 5 6            u w 1       7    1      9
//   u w 7 9
//
// Multilevel hashmap structure:
// {
//   ["s","t"] : {                    <--- group-by field names
//     ["x","y"] : {                  <--- value field names
//       "corr" : stats2_corr object,
//       "cov"  : stats2_cov  object
//     }
//   },
//   ["u","v"] : {
//     ["x","y"] : {
//       "corr" : stats2_corr object,
//       "cov"  : stats2_cov  object
//     }
//   },
//   ["u","w"] : {
//     ["x","y"] : {
//       "corr" : stats2_corr object,
//       "cov"  : stats2_cov  object
//     }
//   },
// }
//
// In the iterative case, add to the current record its current group's stats fields.
// In the non-iterative case, produce output only at the end of the input stream.

func (tr *TransformerStats2) Transform(
	inrecAndContext *types.RecordAndContext,
	outputRecordsAndContexts *[]*types.RecordAndContext, // list of *types.RecordAndContext
	inputDownstreamDoneChannel <-chan bool,
	outputDownstreamDoneChannel chan<- bool,
) {
	HandleDefaultDownstreamDone(inputDownstreamDoneChannel, outputDownstreamDoneChannel)
	if !inrecAndContext.EndOfStream {

		tr.ingest(inrecAndContext)

		if tr.doIterativeStats {
			// The input record is modified in this case, with new fields appended
			*outputRecordsAndContexts = append(*outputRecordsAndContexts, inrecAndContext)
		}
		// if tr.doHoldAndFit, the input record is held by the ingestor

	} else { // end of record stream
		if !tr.doIterativeStats { // in the iterative case, already emitted per-record
			if tr.doHoldAndFit {
				tr.fit(outputRecordsAndContexts)
			} else {
				tr.emit(outputRecordsAndContexts, &inrecAndContext.Context)
			}
		}
		*outputRecordsAndContexts = append(*outputRecordsAndContexts, inrecAndContext) // end-of-stream marker
	}
}

func (tr *TransformerStats2) ingest(
	inrecAndContext *types.RecordAndContext,
) {
	inrec := inrecAndContext.Record

	// E.g. if grouping by "a" and "b", and the current record has a=circle, b=blue,
	// then groupingKey is the string "circle,blue".
	groupingKey, groupByFieldValues, ok := inrec.GetSelectedValuesAndJoined(tr.groupByFieldNameList)
	if !ok {
		return
	}

	tr.groupingKeysToGroupByFieldValues.Put(groupingKey, groupByFieldValues)

	groupToValueFields := tr.namedAccumulators.Get(groupingKey)
	if groupToValueFields == nil {
		groupToValueFields = lib.NewOrderedMap[*lib.OrderedMap[utils.IStats2Accumulator]]()
		tr.namedAccumulators.Put(groupingKey, groupToValueFields)
	}

	if tr.doHoldAndFit { // Retain the input record in memory, for fitting and delivery at end of stream
		groupToRecords := tr.recordGroups.Get(groupingKey)
		if groupToRecords == nil {
			records := []*types.RecordAndContext{}
			groupToRecords = &records
			tr.recordGroups.Put(groupingKey, groupToRecords)
		}
		*groupToRecords = append(*groupToRecords, inrecAndContext)
	}

	// for [["x","y"]]
	n := len(tr.valueFieldNameList)
	for i := 0; i < n; i += 2 {
		valueFieldName1 := tr.valueFieldNameList[i]
		valueFieldName2 := tr.valueFieldNameList[i+1]

		key := valueFieldName1 + stats2KeySeparator + valueFieldName2

		valueFieldsToAccumulator := groupToValueFields.Get(key)
		if valueFieldsToAccumulator == nil {
			valueFieldsToAccumulator = lib.NewOrderedMap[utils.IStats2Accumulator]()
			groupToValueFields.Put(key, valueFieldsToAccumulator)
		}

		mval1 := inrec.Get(valueFieldName1)
		mval2 := inrec.Get(valueFieldName2)
		if mval1 == nil || mval2 == nil { // Key absent in current record
			continue
		}
		if mval1.IsVoid() || mval2.IsVoid() { // Key present in current record but with empty value
			continue
		}

		// for ["corr", "cov"]
		for _, accumulatorName := range tr.accumulatorNameList {
			accumulator := valueFieldsToAccumulator.Get(accumulatorName)
			if accumulator == nil {
				accumulator = tr.accumulatorFactory.Make(
					valueFieldName1,
					valueFieldName2,
					accumulatorName,
					tr.doVerbose,
				)
				if accumulator == nil {
					fmt.Fprintf(os.Stderr, "mlr stats2: accumulator creation failed\n")
					os.Exit(1)
				}
				valueFieldsToAccumulator.Put(accumulatorName, accumulator)
			}
			accumulator.Ingest(
				mval1.GetNumericToFloatValueOrDie(),
				mval2.GetNumericToFloatValueOrDie(),
			)
		}

		if tr.doIterativeStats {
			tr.populateRecord(
				inrecAndContext.Record,
				valueFieldName1,
				valueFieldName2,
				valueFieldsToAccumulator,
			)
		}
	}
}

func (tr *TransformerStats2) emit(
	outputRecordsAndContexts *[]*types.RecordAndContext, // list of *types.RecordAndContext
	context *types.Context,
) {
	for pa := tr.namedAccumulators.Head; pa != nil; pa = pa.Next {
		outrec := mlrval.NewMlrmapAsRecord()

		// Add in a=s,b=t fields:
		groupingKey := pa.Key
		groupByFieldValues := tr.groupingKeysToGroupByFieldValues.Get(groupingKey)
		for i, groupByFieldName := range tr.groupByFieldNameList {
			outrec.PutReference(groupByFieldName, groupByFieldValues[i].Copy())
		}

		// Add in fields such as x_y_corr, etc.
		groupToValueFields := tr.namedAccumulators.Get(groupingKey)

		// For "x","y"
		for pc := groupToValueFields.Head; pc != nil; pc = pc.Next {

			pairs := strings.Split(pc.Key, stats2KeySeparator)
			valueFieldName1 := pairs[0]
			valueFieldName2 := pairs[1]
			valueFieldsToAccumulator := pc.Value

			tr.populateRecord(outrec, valueFieldName1, valueFieldName2, valueFieldsToAccumulator)

			// For "corr", "linreg"
			for pd := valueFieldsToAccumulator.Head; pd != nil; pd = pd.Next {
				accumulator := pd.Value
				accumulator.Populate(valueFieldName1, valueFieldName2, outrec)
			}
		}

		*outputRecordsAndContexts = append(*outputRecordsAndContexts, types.NewRecordAndContext(outrec, context))
	}
}

func (tr *TransformerStats2) populateRecord(
	outrec *mlrval.Mlrmap,
	valueFieldName1 string,
	valueFieldName2 string,
	valueFieldsToAccumulator *lib.OrderedMap[utils.IStats2Accumulator],
) {
	// For "corr", "linreg"
	for pe := valueFieldsToAccumulator.Head; pe != nil; pe = pe.Next {
		accumulator := pe.Value
		accumulator.Populate(valueFieldName1, valueFieldName2, outrec)
	}
}

func (tr *TransformerStats2) fit(
	outputRecordsAndContexts *[]*types.RecordAndContext, // list of *types.RecordAndContext
) {
	for pa := tr.namedAccumulators.Head; pa != nil; pa = pa.Next {
		groupingKey := pa.Key
		groupToValueFields := pa.Value
		recordsAndContexts := tr.recordGroups.Get(groupingKey)
		if recordsAndContexts == nil {
			continue
		}
		for _, recordAndContext := range *recordsAndContexts {
			record := recordAndContext.Record

			// For "x","y"
			for pb := groupToValueFields.Head; pb != nil; pb = pb.Next {
				pairs := strings.Split(pb.Key, stats2KeySeparator)
				valueFieldName1 := pairs[0]
				valueFieldName2 := pairs[1]
				valueFieldsToAccumulator := pb.Value

				// For "linreg-ols", "logireg"
				for pc := valueFieldsToAccumulator.Head; pc != nil; pc = pc.Next {
					accumulator := pc.Value

					// Note R2, cov, corr, etc have no non-trivial fit-function
					mval1 := record.Get(valueFieldName1)
					mval2 := record.Get(valueFieldName2)
					if mval1 != nil && mval2 != nil {
						accumulator.Fit(
							mval1.GetNumericToFloatValueOrDie(),
							mval2.GetNumericToFloatValueOrDie(),
							record,
						)
					}
				}
			}

			*outputRecordsAndContexts = append(*outputRecordsAndContexts, recordAndContext)
		}
		*recordsAndContexts = (*recordsAndContexts)[:0]
	}
}
