src/pages/heatmap/index.tsx

import { useState, useEffect } from 'react';
import ActivityCalendar from 'react-activity-calendar';
import { Tooltip } from '@mui/material';
import Link from 'next/link';
import { fetchAllModelData, generateCalendarData, aggregateCalendarData, CalendarData, Activity } from '@/utils/modelData';
import { PROVIDERS_MAP } from '@/utils/providers';

interface OpenSourceHeatmapProps {
  calendarData: CalendarData & { total: Activity[] };
}

export const getStaticProps = async () => {
  try {
    const modelData = await fetchAllModelData();
    const calendarData = generateCalendarData(modelData);
    
    return {
      props: {
        calendarData: {
          total: aggregateCalendarData(calendarData),
          ...calendarData
        }
      },
      revalidate: process.env.VERCEL_ENV === 'production' ? 43200 : 3600, // 12 hours in production, 1 hour in development
    };
  } catch (error) {
    console.error('Error fetching model data:', error);
    return {
      props: {
        calendarData: { total: [] }
      },
      revalidate: 3600, // 1 hour on error
    };
  }
};

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

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

  return (
    <div className="w-full max-w-screen-lg mx-auto p-4">
      <div className="flex justify-between items-center mb-8">
        <div>
          <h1 className="text-3xl lg:text-4xl font-bold mb-2">Chinese AI Open Source Heatmap 🤗</h1>
          <p className="text-gray-600 dark:text-gray-400">
            Visualize the activity of Chinese AI models and datasets
          </p>
        </div>
        <Link 
          href="/trend"
          className="px-4 py-2 rounded-lg bg-blue-500 text-white hover:bg-blue-600 transition-colors"
        >
          View Trends
        </Link>
      </div>
      {isLoading ? (
        <p className="text-center">Loading...</p>
      ) : (
        <div className="space-y-8 mx-10">
          {/* Total Activity Heatmap */}
          <div className="flex flex-col">
            <h2 className="text-2xl font-bold mb-2">Total Activity</h2>
            <div className="w-full overflow-x-auto">
              <div className="inline-block mx-auto">
                <ActivityCalendar
                  data={calendarData.total}
                  theme={{
                    dark: ['#161b22', '#4CAF50'],
                    light: ['#161b22', '#4CAF50'],
                  }}
                  colorScheme="dark"
                  hideTotalCount
                  renderBlock={(block, activity) => (
                    <Tooltip
                      title={`${activity.count} total models created on ${activity.date}`}
                      arrow
                    >
                      {block}
                    </Tooltip>
                  )}
                />
              </div>
            </div>
          </div>

          <h2 className="text-2xl font-bold mt-12 mb-6">Activity by Organization</h2>
          
          {/* Individual Provider Heatmaps */}
          {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">
                <h2 className="text-xl font-bold mb-2">{providerName}</h2>
                <div className="w-full overflow-x-auto">
                  <div className="inline-block mx-auto">
                    <ActivityCalendar
                      data={calendarData[providerName] || []}
                      theme={{
                        dark: ['#161b22', color],
                        light: ['#161b22', color],
                      }}
                      colorScheme="dark"
                      hideTotalCount
                      renderBlock={(block, activity) => (
                        <Tooltip
                          title={`${activity.count} models created on ${activity.date}`}
                          arrow
                        >
                          {block}
                        </Tooltip>
                      )}
                    />
                  </div>
                </div>
              </div>
            ))}
        </div>
      )}
    </div>
  );
};

export default OpenSourceHeatmap;