gooseek/backend/internal/finance/heatmap.go

package finance

import (
	"context"
	"encoding/json"
	"fmt"
	"math"
	"net/http"
	"sort"
	"strings"
	"sync"
	"time"
)

type HeatmapService struct {
	cache      map[string]*CachedHeatmap
	mu         sync.RWMutex
	httpClient *http.Client
	config     HeatmapConfig
}

type HeatmapConfig struct {
	DataProviderURL string
	CacheTTL        time.Duration
	RefreshInterval time.Duration
}

type CachedHeatmap struct {
	Data      *MarketHeatmap
	ExpiresAt time.Time
}

type MarketHeatmap struct {
	ID          string                `json:"id"`
	Title       string                `json:"title"`
	Type        HeatmapType           `json:"type"`
	Market      string                `json:"market"`
	Sectors     []Sector              `json:"sectors"`
	Tickers     []TickerData          `json:"tickers"`
	Summary     MarketSummary         `json:"summary"`
	UpdatedAt   time.Time             `json:"updatedAt"`
	TimeRange   string                `json:"timeRange"`
	Colorscale  Colorscale            `json:"colorscale"`
}

type HeatmapType string

const (
	HeatmapTreemap     HeatmapType = "treemap"
	HeatmapGrid        HeatmapType = "grid"
	HeatmapBubble      HeatmapType = "bubble"
	HeatmapSectorChart HeatmapType = "sector_chart"
)

type Sector struct {
	ID           string        `json:"id"`
	Name         string        `json:"name"`
	Change       float64       `json:"change"`
	MarketCap    float64       `json:"marketCap"`
	Volume       float64       `json:"volume"`
	TickerCount  int           `json:"tickerCount"`
	TopGainers   []TickerData  `json:"topGainers,omitempty"`
	TopLosers    []TickerData  `json:"topLosers,omitempty"`
	Color        string        `json:"color"`
	Weight       float64       `json:"weight"`
}

type TickerData struct {
	Symbol         string    `json:"symbol"`
	Name           string    `json:"name"`
	Price          float64   `json:"price"`
	Change         float64   `json:"change"`
	ChangePercent  float64   `json:"changePercent"`
	Volume         float64   `json:"volume"`
	MarketCap      float64   `json:"marketCap"`
	Sector         string    `json:"sector"`
	Industry       string    `json:"industry"`
	Color          string    `json:"color"`
	Size           float64   `json:"size"`
	PrevClose      float64   `json:"prevClose,omitempty"`
	DayHigh        float64   `json:"dayHigh,omitempty"`
	DayLow         float64   `json:"dayLow,omitempty"`
	Week52High     float64   `json:"week52High,omitempty"`
	Week52Low      float64   `json:"week52Low,omitempty"`
	PE             float64   `json:"pe,omitempty"`
	EPS            float64   `json:"eps,omitempty"`
	Dividend       float64   `json:"dividend,omitempty"`
	DividendYield  float64   `json:"dividendYield,omitempty"`
}

type MarketSummary struct {
	TotalMarketCap    float64 `json:"totalMarketCap"`
	TotalVolume       float64 `json:"totalVolume"`
	AdvancingCount    int     `json:"advancingCount"`
	DecliningCount    int     `json:"decliningCount"`
	UnchangedCount    int     `json:"unchangedCount"`
	AverageChange     float64 `json:"averageChange"`
	TopGainer         *TickerData `json:"topGainer,omitempty"`
	TopLoser          *TickerData `json:"topLoser,omitempty"`
	MostActive        *TickerData `json:"mostActive,omitempty"`
	MarketSentiment   string  `json:"marketSentiment"`
	VIX               float64 `json:"vix,omitempty"`
	FearGreedIndex    int     `json:"fearGreedIndex,omitempty"`
}

type Colorscale struct {
	Min          float64   `json:"min"`
	Max          float64   `json:"max"`
	MidPoint     float64   `json:"midPoint"`
	Colors       []string  `json:"colors"`
	Thresholds   []float64 `json:"thresholds"`
}

var DefaultColorscale = Colorscale{
	Min:      -10,
	Max:      10,
	MidPoint: 0,
	Colors: []string{
		"#ef4444",
		"#f87171",
		"#fca5a5",
		"#fecaca",
		"#e5e7eb",
		"#bbf7d0",
		"#86efac",
		"#4ade80",
		"#22c55e",
	},
	Thresholds: []float64{-5, -3, -2, -1, 1, 2, 3, 5},
}

func NewHeatmapService(cfg HeatmapConfig) *HeatmapService {
	if cfg.CacheTTL == 0 {
		cfg.CacheTTL = 5 * time.Minute
	}
	if cfg.RefreshInterval == 0 {
		cfg.RefreshInterval = time.Minute
	}

	return &HeatmapService{
		cache:      make(map[string]*CachedHeatmap),
		httpClient: &http.Client{Timeout: 30 * time.Second},
		config:     cfg,
	}
}

func (s *HeatmapService) GetMarketHeatmap(ctx context.Context, market string, timeRange string) (*MarketHeatmap, error) {
	cacheKey := fmt.Sprintf("%s:%s", market, timeRange)

	s.mu.RLock()
	if cached, ok := s.cache[cacheKey]; ok && time.Now().Before(cached.ExpiresAt) {
		s.mu.RUnlock()
		return cached.Data, nil
	}
	s.mu.RUnlock()

	heatmap, err := s.fetchMarketData(ctx, market, timeRange)
	if err != nil {
		return nil, err
	}

	s.mu.Lock()
	s.cache[cacheKey] = &CachedHeatmap{
		Data:      heatmap,
		ExpiresAt: time.Now().Add(s.config.CacheTTL),
	}
	s.mu.Unlock()

	return heatmap, nil
}

func (s *HeatmapService) GetSectorHeatmap(ctx context.Context, market, sector, timeRange string) (*MarketHeatmap, error) {
	heatmap, err := s.GetMarketHeatmap(ctx, market, timeRange)
	if err != nil {
		return nil, err
	}

	filteredTickers := make([]TickerData, 0)
	for _, t := range heatmap.Tickers {
		if strings.EqualFold(t.Sector, sector) {
			filteredTickers = append(filteredTickers, t)
		}
	}

	sectorHeatmap := &MarketHeatmap{
		ID:         fmt.Sprintf("%s-%s", market, sector),
		Title:      fmt.Sprintf("%s - %s", market, sector),
		Type:       HeatmapTreemap,
		Market:     market,
		Tickers:    filteredTickers,
		TimeRange:  timeRange,
		UpdatedAt:  time.Now(),
		Colorscale: DefaultColorscale,
	}

	sectorHeatmap.Summary = s.calculateSummary(filteredTickers)

	return sectorHeatmap, nil
}

func (s *HeatmapService) fetchMarketData(ctx context.Context, market, timeRange string) (*MarketHeatmap, error) {
	heatmap := s.generateMockMarketData(market)
	heatmap.TimeRange = timeRange

	return heatmap, nil
}

func (s *HeatmapService) generateMockMarketData(market string) *MarketHeatmap {
	sectors := []struct {
		name     string
		tickers  []struct{ symbol, name string }
	}{
		{"Technology", []struct{ symbol, name string }{
			{"AAPL", "Apple Inc."},
			{"MSFT", "Microsoft Corp."},
			{"GOOGL", "Alphabet Inc."},
			{"AMZN", "Amazon.com Inc."},
			{"META", "Meta Platforms"},
			{"NVDA", "NVIDIA Corp."},
			{"TSLA", "Tesla Inc."},
		}},
		{"Healthcare", []struct{ symbol, name string }{
			{"JNJ", "Johnson & Johnson"},
			{"UNH", "UnitedHealth Group"},
			{"PFE", "Pfizer Inc."},
			{"MRK", "Merck & Co."},
			{"ABBV", "AbbVie Inc."},
		}},
		{"Finance", []struct{ symbol, name string }{
			{"JPM", "JPMorgan Chase"},
			{"BAC", "Bank of America"},
			{"WFC", "Wells Fargo"},
			{"GS", "Goldman Sachs"},
			{"MS", "Morgan Stanley"},
		}},
		{"Energy", []struct{ symbol, name string }{
			{"XOM", "Exxon Mobil"},
			{"CVX", "Chevron Corp."},
			{"COP", "ConocoPhillips"},
			{"SLB", "Schlumberger"},
		}},
		{"Consumer", []struct{ symbol, name string }{
			{"WMT", "Walmart Inc."},
			{"PG", "Procter & Gamble"},
			{"KO", "Coca-Cola Co."},
			{"PEP", "PepsiCo Inc."},
			{"COST", "Costco Wholesale"},
		}},
	}

	allTickers := make([]TickerData, 0)
	allSectors := make([]Sector, 0)

	for _, sec := range sectors {
		sectorTickers := make([]TickerData, 0)
		sectorChange := 0.0

		for _, t := range sec.tickers {
			change := (randomFloat(-5, 5))
			price := randomFloat(50, 500)
			marketCap := randomFloat(50e9, 3000e9)
			volume := randomFloat(1e6, 100e6)

			ticker := TickerData{
				Symbol:        t.symbol,
				Name:          t.name,
				Price:         price,
				Change:        price * change / 100,
				ChangePercent: change,
				Volume:        volume,
				MarketCap:     marketCap,
				Sector:        sec.name,
				Color:         getColorForChange(change),
				Size:          math.Log10(marketCap) * 10,
			}

			sectorTickers = append(sectorTickers, ticker)
			sectorChange += change
		}

		if len(sectorTickers) > 0 {
			sectorChange /= float64(len(sectorTickers))
		}

		sort.Slice(sectorTickers, func(i, j int) bool {
			return sectorTickers[i].ChangePercent > sectorTickers[j].ChangePercent
		})

		var topGainers, topLosers []TickerData
		if len(sectorTickers) >= 2 {
			topGainers = sectorTickers[:2]
			topLosers = sectorTickers[len(sectorTickers)-2:]
		}

		sectorMarketCap := 0.0
		sectorVolume := 0.0
		for _, t := range sectorTickers {
			sectorMarketCap += t.MarketCap
			sectorVolume += t.Volume
		}

		sector := Sector{
			ID:          strings.ToLower(strings.ReplaceAll(sec.name, " ", "_")),
			Name:        sec.name,
			Change:      sectorChange,
			MarketCap:   sectorMarketCap,
			Volume:      sectorVolume,
			TickerCount: len(sectorTickers),
			TopGainers:  topGainers,
			TopLosers:   topLosers,
			Color:       getColorForChange(sectorChange),
			Weight:      sectorMarketCap,
		}

		allSectors = append(allSectors, sector)
		allTickers = append(allTickers, sectorTickers...)
	}

	sort.Slice(allTickers, func(i, j int) bool {
		return allTickers[i].MarketCap > allTickers[j].MarketCap
	})

	return &MarketHeatmap{
		ID:         market,
		Title:      getMarketTitle(market),
		Type:       HeatmapTreemap,
		Market:     market,
		Sectors:    allSectors,
		Tickers:    allTickers,
		Summary:    *s.calculateSummaryPtr(allTickers),
		UpdatedAt:  time.Now(),
		Colorscale: DefaultColorscale,
	}
}

func (s *HeatmapService) calculateSummary(tickers []TickerData) MarketSummary {
	return *s.calculateSummaryPtr(tickers)
}

func (s *HeatmapService) calculateSummaryPtr(tickers []TickerData) *MarketSummary {
	summary := &MarketSummary{}

	var totalChange float64
	var topGainer, topLoser, mostActive *TickerData

	for i := range tickers {
		t := &tickers[i]
		summary.TotalMarketCap += t.MarketCap
		summary.TotalVolume += t.Volume
		totalChange += t.ChangePercent

		if t.ChangePercent > 0 {
			summary.AdvancingCount++
		} else if t.ChangePercent < 0 {
			summary.DecliningCount++
		} else {
			summary.UnchangedCount++
		}

		if topGainer == nil || t.ChangePercent > topGainer.ChangePercent {
			topGainer = t
		}
		if topLoser == nil || t.ChangePercent < topLoser.ChangePercent {
			topLoser = t
		}
		if mostActive == nil || t.Volume > mostActive.Volume {
			mostActive = t
		}
	}

	if len(tickers) > 0 {
		summary.AverageChange = totalChange / float64(len(tickers))
	}

	summary.TopGainer = topGainer
	summary.TopLoser = topLoser
	summary.MostActive = mostActive

	if summary.AverageChange > 1 {
		summary.MarketSentiment = "bullish"
	} else if summary.AverageChange < -1 {
		summary.MarketSentiment = "bearish"
	} else {
		summary.MarketSentiment = "neutral"
	}

	return summary
}

func (s *HeatmapService) GenerateTreemapData(heatmap *MarketHeatmap) interface{} {
	children := make([]map[string]interface{}, 0)

	for _, sector := range heatmap.Sectors {
		sectorChildren := make([]map[string]interface{}, 0)

		for _, ticker := range heatmap.Tickers {
			if ticker.Sector == sector.Name {
				sectorChildren = append(sectorChildren, map[string]interface{}{
					"name":   ticker.Symbol,
					"value":  ticker.MarketCap,
					"change": ticker.ChangePercent,
					"color":  ticker.Color,
					"data":   ticker,
				})
			}
		}

		children = append(children, map[string]interface{}{
			"name":     sector.Name,
			"children": sectorChildren,
			"change":   sector.Change,
			"color":    sector.Color,
		})
	}

	return map[string]interface{}{
		"name":     heatmap.Market,
		"children": children,
	}
}

func (s *HeatmapService) GenerateGridData(heatmap *MarketHeatmap, rows, cols int) [][]TickerData {
	grid := make([][]TickerData, rows)
	for i := range grid {
		grid[i] = make([]TickerData, cols)
	}

	idx := 0
	for i := 0; i < rows && idx < len(heatmap.Tickers); i++ {
		for j := 0; j < cols && idx < len(heatmap.Tickers); j++ {
			grid[i][j] = heatmap.Tickers[idx]
			idx++
		}
	}

	return grid
}

func (s *HeatmapService) GetTopMovers(ctx context.Context, market string, count int) (*TopMovers, error) {
	heatmap, err := s.GetMarketHeatmap(ctx, market, "1d")
	if err != nil {
		return nil, err
	}

	tickers := make([]TickerData, len(heatmap.Tickers))
	copy(tickers, heatmap.Tickers)

	sort.Slice(tickers, func(i, j int) bool {
		return tickers[i].ChangePercent > tickers[j].ChangePercent
	})

	gainers := tickers
	if len(gainers) > count {
		gainers = gainers[:count]
	}

	sort.Slice(tickers, func(i, j int) bool {
		return tickers[i].ChangePercent < tickers[j].ChangePercent
	})

	losers := tickers
	if len(losers) > count {
		losers = losers[:count]
	}

	sort.Slice(tickers, func(i, j int) bool {
		return tickers[i].Volume > tickers[j].Volume
	})

	active := tickers
	if len(active) > count {
		active = active[:count]
	}

	return &TopMovers{
		Gainers:    gainers,
		Losers:     losers,
		MostActive: active,
		UpdatedAt:  time.Now(),
	}, nil
}

type TopMovers struct {
	Gainers    []TickerData `json:"gainers"`
	Losers     []TickerData `json:"losers"`
	MostActive []TickerData `json:"mostActive"`
	UpdatedAt  time.Time    `json:"updatedAt"`
}

func getColorForChange(change float64) string {
	if change >= 5 {
		return "#22c55e"
	} else if change >= 3 {
		return "#4ade80"
	} else if change >= 1 {
		return "#86efac"
	} else if change >= 0 {
		return "#bbf7d0"
	} else if change >= -1 {
		return "#fecaca"
	} else if change >= -3 {
		return "#fca5a5"
	} else if change >= -5 {
		return "#f87171"
	}
	return "#ef4444"
}

func getMarketTitle(market string) string {
	titles := map[string]string{
		"sp500":    "S&P 500",
		"nasdaq":   "NASDAQ",
		"dow":      "Dow Jones",
		"moex":     "MOEX",
		"crypto":   "Cryptocurrency",
		"forex":    "Forex",
		"commodities": "Commodities",
	}
	if title, ok := titles[strings.ToLower(market)]; ok {
		return title
	}
	return market
}

var rng uint64 = uint64(time.Now().UnixNano())

func randomFloat(min, max float64) float64 {
	rng ^= rng << 13
	rng ^= rng >> 17
	rng ^= rng << 5
	f := float64(rng) / float64(1<<64)
	return min + f*(max-min)
}

func (h *MarketHeatmap) ToJSON() ([]byte, error) {
	return json.Marshal(h)
}