"use client";

import { Inter } from "next/font/google";
import ActivityCalendar from "react-activity-calendar";
import { useState, useEffect } from "react";
import { Tooltip as MuiTooltip } from '@mui/material';

const inter = Inter({ subsets: ["latin"] });

interface ModelData {
  createdAt: string;
  id: string;
}

interface Activity {
  date: string;
  count: number;
  level: number;
}

export default function Home() {
  const [calendarData, setCalendarData] = useState<Record<string, Activity[]>>({});
  const [isLoading, setIsLoading] = useState(true);

  const PROVIDERS_MAP: Record<string, { color: string; authors: string[] }> = {
    "BAAI": { color: "#ff7000", authors: ["BAAI"] },        // Vibrant Orange
    "DeepSeek": { color: "#007ACC", authors: ["deepseek-ai"] }, // Light Blue
    "Shanghai AI Lab": { color: "#10A37F", authors: ["internlm", "OpenGVLab", "openmmlab"] }, // Fresh Green
    "Alibaba": { color: "#D2691E", authors: ["Qwen", "Alibaba-NLP", "alibaba-pai", "DAMO-NLP-SG"] },        // Chocolate
    "GLM": { color: "#4285F4", authors: ["THUDM"] },        // Classic Blue
    "Tencent": { color: "#1DA1F2", authors: ["TencentARC", "Tencent-Hunyuan"] }, // Twitter Blue
    "Yi/01": { color: "#FF4500", authors: ["01-ai"] },      // Orange-Red
  }

  const generateCalendarData = (modelData: ModelData[]) => {
    const data: Record<string, Activity[]> = Object.fromEntries(
      Object.keys(PROVIDERS_MAP).map(provider => [provider, []])
    );

    const today = new Date();
    const startDate = new Date(today);
    startDate.setMonth(today.getMonth() - 11);
    startDate.setDate(1); // start from the first day of the month

    // generate daily data for each provider
    for (let d = new Date(startDate); d <= today; d.setDate(d.getDate() + 1)) {
      const dateString = d.toISOString().split('T')[0];
      
      Object.entries(PROVIDERS_MAP).forEach(([provider, { authors }]) => {
        const count = modelData.filter(item => 
          item.createdAt.startsWith(dateString) && 
          authors.some(author => item.id.startsWith(author))
        ).length;
        
        data[provider].push({ date: dateString, count, level: 0 });
      });
    }

    // calculate average counts for each provider
    const avgCounts = Object.fromEntries(
      Object.entries(data).map(([provider, days]) => [
        provider,
        days.reduce((sum, day) => sum + day.count, 0) / days.length || 0
      ])
    );

    // assign levels based on count relative to average
    Object.entries(data).forEach(([provider, days]) => {
      const avgCount = avgCounts[provider];
      days.forEach(day => {
        day.level = 
          day.count === 0 ? 0 :
          day.count <= avgCount * 0.5 ? 1 :
          day.count <= avgCount ? 2 :
          day.count <= avgCount * 1.5 ? 3 : 4;
      });
    });

    return data;
  }

  const initData = async () => {
    try {
      const allAuthors = Object.values(PROVIDERS_MAP).flatMap(({ authors }) => authors);
      const uniqueAuthors = Array.from(new Set(allAuthors));

      const allModelData = await Promise.all(
        uniqueAuthors.map(async (author) => {
          const response = await fetch(`https://huggingface.co/api/models?author=${author}&sort=createdAt&direction=-1`);
          const data = await response.json();
          return data.map((item: any) => ({
            createdAt: item.createdAt,
            id: item.id,
          }));
        })
      );

      const flatModelData = allModelData.flat();
      const calendarData = generateCalendarData(flatModelData);
      setCalendarData(calendarData);
    } catch (error) {
      console.error("Error fetching data:", error);
    } finally {
      setIsLoading(false);
    }
  }

  useEffect(() => {
    initData();
  }, []);

  return (
    <main className={`flex flex-col items-center justify-center min-h-screen mx-auto p-24 ${inter.className}`}>
      <h1 className="text-5xl font-bold text-center">Open Source Heatmap</h1>
      <p className="text-center mt-2 text-sm">A heatmap for open source model releases.</p>
      <div className="mt-16">
        {isLoading ? (
          <p>Loading...</p>
        ) : (
          <>
            {Object.entries(PROVIDERS_MAP)
              .sort(([keyA], [keyB]) => 
                calendarData[keyB].reduce((sum, day) => sum + day.count, 0) -
                calendarData[keyA].reduce((sum, day) => sum + day.count, 0)
              )
              .map(([providerName, { color }]) => (
                <div key={providerName} className="mb-8">
                  <h2 className="text-2xl font-bold mb-2">{providerName}</h2>
                  <ActivityCalendar 
                    data={calendarData[providerName]}
                    theme={{
                      dark: ['#161b22', color],
                      light: ['#e0e0e0', color],
                    }}
                    hideTotalCount
                    renderBlock={(block, activity) => (
                      <MuiTooltip
                        title={`${activity.count} models created on ${activity.date}`}
                      >
                        {block}
                      </MuiTooltip>
                    )}
                  />
                </div>
              ))
            }
          </>
        )}
      </div>
    </main>
  );
}