src/pages/index.tsx

import React, { useState, useEffect } from 'react';
import ActivityCalendar from "react-activity-calendar";
import { Tooltip } from '@mui/material';

interface ProviderInfo {
  color: string;
  authors: string[];
}

const PROVIDERS_MAP: Record<string, ProviderInfo> = {
  "Mistral AI": { color: "#ff7000", authors: ["mistralai"] },
  "Meta": { color: "#1877F2", authors: ["facebook", "meta-llama"] },
  "OpenAI": { color: "#10A37F", authors: ["openai"] },
  "Anthropic": { color: "#cc785c", authors: ["Anthropic"] },
  "Google": { color: "#DB4437", authors: ["google"] },
  "Allen Institute for AI": { color: "#5E35B1", authors: ["allenai"] },
  "Apple": { color: "#0088cc", authors: ["apple"] },
  "Microsoft": { color: "#FEB800", authors: ["microsoft"] },
  "NVIDIA": { color: "#76B900", authors: ["nvidia"] },
  "DeepSeek": { color: "#0088cc", authors: ["deepseek-ai"] },
  "Qwen": { color: "#0088cc", authors: ["Qwen"] },
  "Cohere For AI": { color: "#4C6EE6", authors: ["CohereForAI"] },
  "IBM": { color: "#4C6EE6", authors: ["ibm-granite"] },
  "Stability AI": { color: "#A020F0", authors: ["stabilityai"] },
};

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

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

interface CalendarData {
  [key: string]: Activity[];
}

const generateCalendarData = (modelData: ModelData[]): CalendarData => {
  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);

  // create a map to store counts for each provider and date
  const countMap: Record<string, Record<string, number>> = {};

  modelData.forEach(item => {
    const dateString = item.createdAt.split('T')[0];
    Object.entries(PROVIDERS_MAP).forEach(([provider, { authors }]) => {
      if (authors.some(author => item.id.startsWith(author))) {
        countMap[provider] = countMap[provider] || {};
        countMap[provider][dateString] = (countMap[provider][dateString] || 0) + 1;
      }
    });
  });

  // fill in with 0s for days with no activity
  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]) => {
      const count = countMap[provider]?.[dateString] || 0;
      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;
}

export async function getStaticProps() {
  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);

    return {
      props: {
        calendarData,
      },
      revalidate: 3600, // regenerate every hour
    };
  } catch (error) {
    console.error("Error fetching data:", error);
    return {
      props: {
        calendarData: {},
      },
      revalidate: 60, // retry after 1 minute if there was an error
    };
  }
}

interface OpenSourceHeatmapProps {
  calendarData: CalendarData;
}

const OpenSourceHeatmap: React.FC<OpenSourceHeatmapProps> = ({ calendarData }) => {
  const [isLoading, setIsLoading] = useState(true);

  useEffect(() => {
    if (calendarData && Object.keys(calendarData).length > 0) {
      setIsLoading(false);
    }
  }, [calendarData]);

  return (
    <div className="w-full max-w-screen-lg mx-auto p-4">
      <h1 className="text-3xl lg:text-5xl mt-16 font-bold text-center mb-2">Hugging Face Heatmap 🤗</h1>
      <p className="text-center text-sm mb-8">
        The top AI labs and model releases. <br />
        Request more heatmaps by{' '}
        <a
          href="https://huggingface.co/spaces/cfahlgren1/model-release-heatmap/discussions/new"
          target="_blank"
          rel="noopener noreferrer"
          className="text-blue-500 hover:underline"
        >
          opening a discussion
        </a>
        .
      </p>
      {isLoading ? (
        <p className="text-center">Loading...</p>
      ) : (
        <div className="space-y-8">
          {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="flex flex-col items-center">
                <h2 className="text-xl font-bold mb-4">{providerName}</h2>
                <div className="w-full overflow-x-auto flex justify-center">
                  <ActivityCalendar 
                    data={calendarData[providerName]}
                    theme={{
                      dark: ['#161b22', color],
                      light: ['#e0e0e0', color],
                    }}
                    hideTotalCount
                    renderBlock={(block, activity) => (
                      <Tooltip
                        title={`${activity.count} models created on ${activity.date}`}
                        arrow
                      >
                        {block}
                      </Tooltip>
                    )}
                  />
                </div>
              </div>
            ))
          }
        </div>
      )}
    </div>
  );
}

export default OpenSourceHeatmap;