gitstack/graph/

colors.rs

//! グラフ色管理

use std::collections::{HashMap, VecDeque};

use ratatui::style::Color;

/// メインブランチの色インデックス（予約）
pub const MAIN_BRANCH_COLOR: usize = 0;

/// 色パレットのサイズ
pub const PALETTE_SIZE: usize = 16;

/// グラフで使用する色パレット（16色）
/// 多数の並列ブランチでも識別しやすいよう拡張
const GRAPH_PALETTE: [Color; 16] = [
    Color::Blue,               // 0: メインブランチ用
    Color::Green,              // 1
    Color::Yellow,             // 2
    Color::Magenta,            // 3
    Color::Cyan,               // 4
    Color::Red,                // 5
    Color::LightBlue,          // 6
    Color::LightGreen,         // 7
    Color::LightYellow,        // 8
    Color::LightMagenta,       // 9
    Color::LightCyan,          // 10
    Color::LightRed,           // 11
    Color::Rgb(255, 165, 0),   // 12: Orange
    Color::Rgb(128, 0, 128),   // 13: Purple
    Color::Rgb(0, 128, 128),   // 14: Teal
    Color::Rgb(255, 105, 180), // 15: HotPink
];

/// 色インデックスからColor値を取得
pub fn get_graph_color(idx: usize) -> Color {
    GRAPH_PALETTE[idx % GRAPH_PALETTE.len()]
}

/// 色割り当て管理
pub struct ColorAssigner {
    /// 次に割り当てる色インデックス（0はメインブランチ用に予約）
    next_idx: usize,
    /// 使用中の色インデックス
    used_colors: Vec<bool>,
}

impl ColorAssigner {
    pub fn new() -> Self {
        let mut used_colors = vec![false; GRAPH_PALETTE.len()];
        used_colors[MAIN_BRANCH_COLOR] = true; // メインブランチ色を予約
        Self {
            next_idx: 1,
            used_colors,
        }
    }

    /// メインブランチ用の色を取得
    pub fn main_color(&self) -> usize {
        MAIN_BRANCH_COLOR
    }

    /// 新しい色を割り当て
    pub fn assign(&mut self) -> usize {
        // 空いている色を探す
        for i in 1..GRAPH_PALETTE.len() {
            let idx = (self.next_idx + i - 1) % GRAPH_PALETTE.len();
            if idx != MAIN_BRANCH_COLOR && !self.used_colors[idx] {
                self.used_colors[idx] = true;
                self.next_idx = (idx + 1) % GRAPH_PALETTE.len();
                return idx;
            }
        }
        // 全て使用中なら順番に割り当て
        let idx = if self.next_idx == MAIN_BRANCH_COLOR {
            1
        } else {
            self.next_idx
        };
        self.next_idx = (idx + 1) % GRAPH_PALETTE.len();
        if self.next_idx == MAIN_BRANCH_COLOR {
            self.next_idx = 1;
        }
        idx
    }

    /// 色を解放
    pub fn release(&mut self, idx: usize) {
        if idx != MAIN_BRANCH_COLOR && idx < self.used_colors.len() {
            self.used_colors[idx] = false;
        }
    }
}

impl Default for ColorAssigner {
    fn default() -> Self {
        Self::new()
    }
}

/// 色割り当てのためのコンテキスト情報
#[derive(Debug, Clone)]
pub struct ColorContext {
    /// レーン（列）番号
    pub lane: usize,
    /// メインブランチかどうか
    pub is_main_branch: bool,
    /// 分岐元コミットのハッシュ（同じ分岐元を持つブランチは異なる色にする）
    pub fork_point: Option<String>,
    /// 親コミットの色（継続性のため）
    pub parent_color: Option<usize>,
}

impl ColorContext {
    /// 新しいColorContextを作成
    pub fn new(lane: usize) -> Self {
        Self {
            lane,
            is_main_branch: false,
            fork_point: None,
            parent_color: None,
        }
    }

    /// メインブランチフラグを設定
    pub fn with_main_branch(mut self, is_main: bool) -> Self {
        self.is_main_branch = is_main;
        self
    }

    /// 分岐元を設定
    pub fn with_fork_point(mut self, fork_point: Option<String>) -> Self {
        self.fork_point = fork_point;
        self
    }

    /// 親の色を設定
    pub fn with_parent_color(mut self, parent_color: Option<usize>) -> Self {
        self.parent_color = parent_color;
        self
    }
}

/// ペナルティベースの色割り当て管理
///
/// keifuのアルゴリズムを参考に、以下のペナルティを考慮して色を割り当てる:
/// 1. 隣接レーンペナルティ: 隣接するレーンと同じ色を避ける
/// 2. 履歴ペナルティ: 最近使用した色を避ける
/// 3. Fork siblingsペナルティ: 同じ分岐元を持つブランチと同じ色を避ける
pub struct PenaltyBasedColorAssigner {
    /// 各レーンに割り当てられた色
    lane_colors: HashMap<usize, usize>,
    /// 最近使用した色の履歴
    color_history: VecDeque<usize>,
    /// 分岐元ハッシュ→使用中の色リスト
    fork_colors: HashMap<String, Vec<usize>>,
    /// 履歴サイズ
    history_size: usize,
}

impl PenaltyBasedColorAssigner {
    /// 新しいPenaltyBasedColorAssignerを作成
    pub fn new() -> Self {
        Self {
            lane_colors: HashMap::new(),
            color_history: VecDeque::new(),
            fork_colors: HashMap::new(),
            history_size: 8,
        }
    }

    /// メインブランチ用の色を取得
    pub fn main_color(&self) -> usize {
        MAIN_BRANCH_COLOR
    }

    /// コンテキストを考慮して色を割り当て
    pub fn assign_with_context(&mut self, ctx: &ColorContext) -> usize {
        // メインブランチは常に色0
        if ctx.is_main_branch {
            self.update_state(ctx.lane, MAIN_BRANCH_COLOR, ctx.fork_point.as_ref());
            return MAIN_BRANCH_COLOR;
        }

        // ペナルティ配列を初期化
        let mut penalties = [0.0f32; PALETTE_SIZE];

        // 1. 隣接レーンペナルティ (weight: 10.0)
        //    隣接レーンと同じ色は見分けにくいので高ペナルティ
        for neighbor in [ctx.lane.saturating_sub(1), ctx.lane + 1] {
            if neighbor != ctx.lane {
                if let Some(&color) = self.lane_colors.get(&neighbor) {
                    penalties[color] += 10.0;
                    // 類似色にも軽いペナルティ
                    if let Some(similar) = self.similar_color(color) {
                        penalties[similar] += 5.0;
                    }
                }
            }
        }

        // 2. 履歴ペナルティ (weight: 3.0 * decay)
        //    最近使った色は避ける（垂直方向の重複回避）
        for (age, &color) in self.color_history.iter().enumerate() {
            let decay = 1.0 - (age as f32 / self.history_size as f32);
            penalties[color] += 3.0 * decay;
        }

        // 3. Fork siblingsペナルティ (weight: 8.0)
        //    同じ分岐元から派生したブランチは異なる色にする
        if let Some(ref fork) = ctx.fork_point {
            if let Some(siblings) = self.fork_colors.get(fork) {
                for &color in siblings {
                    penalties[color] += 8.0;
                }
            }
        }

        // 4. メインブランチ色（0）には高ペナルティ
        penalties[MAIN_BRANCH_COLOR] += 100.0;

        // 最小ペナルティの色を選択
        let best = (1..PALETTE_SIZE)
            .min_by(|&a, &b| {
                penalties[a]
                    .partial_cmp(&penalties[b])
                    .unwrap_or(std::cmp::Ordering::Equal)
            })
            .unwrap_or(1);

        // 状態を更新
        self.update_state(ctx.lane, best, ctx.fork_point.as_ref());

        best
    }

    /// 状態を更新
    fn update_state(&mut self, lane: usize, color: usize, fork_point: Option<&String>) {
        self.lane_colors.insert(lane, color);

        self.color_history.push_front(color);
        if self.color_history.len() > self.history_size {
            self.color_history.pop_back();
        }

        if let Some(fork) = fork_point {
            self.fork_colors
                .entry(fork.clone())
                .or_default()
                .push(color);
        }
    }

    /// 類似色を取得（同系統の明暗バリエーション）
    fn similar_color(&self, color: usize) -> Option<usize> {
        // パレット構成:
        // 0: Blue, 1: Green, 2: Yellow, 3: Magenta, 4: Cyan, 5: Red
        // 6: LightBlue, 7: LightGreen, 8: LightYellow, 9: LightMagenta, 10: LightCyan, 11: LightRed
        // 12: Orange, 13: Purple, 14: Teal, 15: HotPink
        match color {
            0 => Some(6), // Blue <-> LightBlue
            6 => Some(0),
            1 => Some(7), // Green <-> LightGreen
            7 => Some(1),
            2 => Some(8), // Yellow <-> LightYellow
            8 => Some(2),
            3 => Some(9), // Magenta <-> LightMagenta
            9 => Some(3),
            4 => Some(10), // Cyan <-> LightCyan
            10 => Some(4),
            5 => Some(11), // Red <-> LightRed
            11 => Some(5),
            _ => None,
        }
    }

    /// レーンの色を取得
    pub fn get_lane_color(&self, lane: usize) -> Option<usize> {
        self.lane_colors.get(&lane).copied()
    }

    /// レーンを解放
    pub fn release_lane(&mut self, lane: usize) {
        self.lane_colors.remove(&lane);
    }
}

impl Default for PenaltyBasedColorAssigner {
    fn default() -> Self {
        Self::new()
    }
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn test_get_graph_color() {
        assert_eq!(get_graph_color(0), Color::Blue);
        assert_eq!(get_graph_color(1), Color::Green);
        assert_eq!(get_graph_color(16), Color::Blue); // ラップアラウンド
    }

    #[test]
    fn test_color_assigner_main_color() {
        let assigner = ColorAssigner::new();
        assert_eq!(assigner.main_color(), MAIN_BRANCH_COLOR);
    }

    #[test]
    fn test_color_assigner_assign_skips_main() {
        let mut assigner = ColorAssigner::new();
        let color1 = assigner.assign();
        assert_ne!(color1, MAIN_BRANCH_COLOR);
    }

    #[test]
    fn test_color_assigner_assign_unique() {
        let mut assigner = ColorAssigner::new();
        let c1 = assigner.assign();
        let c2 = assigner.assign();
        let c3 = assigner.assign();
        assert_ne!(c1, c2);
        assert_ne!(c2, c3);
        assert_ne!(c1, c3);
    }

    #[test]
    fn test_color_assigner_release() {
        let mut assigner = ColorAssigner::new();
        let c1 = assigner.assign();
        let _c2 = assigner.assign();
        assigner.release(c1);
        // 解放後は再利用可能（次回割り当て時に再利用される）
        let c3 = assigner.assign();
        // c1が解放されているので再利用されるはず
        assert!(!assigner.used_colors[c1] || c3 == c1);
    }

    // PenaltyBasedColorAssigner tests

    #[test]
    fn test_penalty_assigner_main_color() {
        let assigner = PenaltyBasedColorAssigner::new();
        assert_eq!(assigner.main_color(), MAIN_BRANCH_COLOR);
    }

    #[test]
    fn test_penalty_assigner_main_branch_always_color_0() {
        let mut assigner = PenaltyBasedColorAssigner::new();
        let ctx = ColorContext::new(0).with_main_branch(true);
        let color = assigner.assign_with_context(&ctx);
        assert_eq!(color, MAIN_BRANCH_COLOR);
    }

    #[test]
    fn test_penalty_assigner_skips_main_color() {
        let mut assigner = PenaltyBasedColorAssigner::new();
        let ctx = ColorContext::new(1);
        let color = assigner.assign_with_context(&ctx);
        assert_ne!(color, MAIN_BRANCH_COLOR);
    }

    #[test]
    fn test_penalty_assigner_adjacent_lanes_different_colors() {
        let mut assigner = PenaltyBasedColorAssigner::new();

        // レーン0に色を割り当て
        let ctx0 = ColorContext::new(0);
        let color0 = assigner.assign_with_context(&ctx0);

        // レーン1は隣接するので異なる色になるはず
        let ctx1 = ColorContext::new(1);
        let color1 = assigner.assign_with_context(&ctx1);

        assert_ne!(color0, color1, "隣接レーンは異なる色になるべき");
    }

    #[test]
    fn test_penalty_assigner_fork_siblings_different_colors() {
        let mut assigner = PenaltyBasedColorAssigner::new();
        let fork_point = "base_commit".to_string();

        // 同じ分岐元から派生したブランチ
        let ctx1 = ColorContext::new(1).with_fork_point(Some(fork_point.clone()));
        let color1 = assigner.assign_with_context(&ctx1);

        let ctx2 = ColorContext::new(2).with_fork_point(Some(fork_point.clone()));
        let color2 = assigner.assign_with_context(&ctx2);

        let ctx3 = ColorContext::new(3).with_fork_point(Some(fork_point));
        let color3 = assigner.assign_with_context(&ctx3);

        assert_ne!(color1, color2, "Fork siblingsは異なる色になるべき");
        assert_ne!(color2, color3, "Fork siblingsは異なる色になるべき");
        assert_ne!(color1, color3, "Fork siblingsは異なる色になるべき");
    }

    #[test]
    fn test_penalty_assigner_history_avoidance() {
        let mut assigner = PenaltyBasedColorAssigner::new();

        // 複数の色を連続で割り当て
        let mut colors = Vec::new();
        for i in 0..5 {
            let ctx = ColorContext::new(i * 3); // 離れたレーンを使用
            let color = assigner.assign_with_context(&ctx);
            colors.push(color);
        }

        // 最初の数色は重複しないはず（履歴ペナルティにより）
        for i in 0..4 {
            for j in (i + 1)..5 {
                if i + 1 < j {
                    // 連続していない場合は重複しにくい
                    // （完全な重複回避は保証されないが、傾向として確認）
                }
            }
        }
        // 少なくとも全部同じ色にはならない
        let unique_count = colors
            .iter()
            .collect::<std::collections::HashSet<_>>()
            .len();
        assert!(unique_count >= 3, "履歴により色のバリエーションがあるべき");
    }

    #[test]
    fn test_penalty_assigner_get_lane_color() {
        let mut assigner = PenaltyBasedColorAssigner::new();
        let ctx = ColorContext::new(5);
        let color = assigner.assign_with_context(&ctx);

        assert_eq!(assigner.get_lane_color(5), Some(color));
        assert_eq!(assigner.get_lane_color(0), None);
    }

    #[test]
    fn test_penalty_assigner_release_lane() {
        let mut assigner = PenaltyBasedColorAssigner::new();
        let ctx = ColorContext::new(3);
        let _color = assigner.assign_with_context(&ctx);

        assert!(assigner.get_lane_color(3).is_some());
        assigner.release_lane(3);
        assert!(assigner.get_lane_color(3).is_none());
    }

    #[test]
    fn test_color_context_builder() {
        let ctx = ColorContext::new(5)
            .with_main_branch(true)
            .with_fork_point(Some("abc123".to_string()))
            .with_parent_color(Some(3));

        assert_eq!(ctx.lane, 5);
        assert!(ctx.is_main_branch);
        assert_eq!(ctx.fork_point, Some("abc123".to_string()));
        assert_eq!(ctx.parent_color, Some(3));
    }

    #[test]
    fn test_penalty_assigner_many_lanes() {
        let mut assigner = PenaltyBasedColorAssigner::new();

        // 16レーン以上を割り当て
        for i in 0..20 {
            let ctx = ColorContext::new(i);
            let color = assigner.assign_with_context(&ctx);
            // 色はパレットサイズ内に収まる
            assert!(color < PALETTE_SIZE, "色インデックスがパレットサイズ内");
        }
    }
}