package main

import (
	"fmt"
	"math"
	"sort"
)

func rank(data []float64) []float64 {
	n := len(data)
	ranked := make([]float64, n)
	for i := range ranked {
		ranked[i] = float64(i + 1)
	}
	sort.Float64s(data)

	for i := 0; i < n; {
		j := i
		for j < n && data[i] == data[j] {
			j++
		}
		rank := (float64(i+1) + float64(j)) / 2
		for k := i; k < j; k++ {
			ranked[k] = rank
		}
		i = j
	}

	return ranked
}

func mannWhitneyU(sampleA, sampleB []float64) (float64, float64) {
	n1 := float64(len(sampleA))
	n2 := float64(len(sampleB))

	all := append(sampleA, sampleB...)
	ranks := rank(all)

	rankSumA := 0.0
	for i := 0; i < len(sampleA); i++ {
		rankSumA += ranks[i]
	}

	u1 := rankSumA - n1*(n1+1)/2
	u2 := n1*n2 - u1

	return u1, u2
}

func calculatePValue(u, n1, n2 float64) float64 {
	meanU := n1 * n2 / 2
	sdU := math.Sqrt(n1 * n2 * (n1 + n2 + 1) / 12)
	z := (u - meanU) / sdU
	pValue := 2 * (1 - normalCDF(math.Abs(z)))
	return pValue
}

func normalCDF(x float64) float64 {
	return 0.5 * (1 + math.Erf(x/math.Sqrt2))
}

func main() {
	sampleA := []float64{-30, -25, -20, -15, -10, -5, 0, 5, 10, 15, 20, 25, 30, 35, 40, 45, 50, 55, 60, 65}
	sampleB := []float64{-28, -23, -18, -13, -8, -3, 2, 7, 12, 17, 22, 27, 32, 37, 42, 47, 52, 57, 62, 67}

	u1, u2 := mannWhitneyU(sampleA, sampleB)
	pValue := calculatePValue(u1, float64(len(sampleA)), float64(len(sampleB)))

	fmt.Printf("U1 统计量: %.2f\n", u1)
	fmt.Printf("U2 统计量: %.2f\n", u2)
	fmt.Printf("p 值: %.4f\n", pValue)
}