import { ModelData, ProviderInfo, CalendarData, Activity } from "../types/heatmap";

export const generateCalendarData = (
    modelData: ModelData[],
    providers: Record<string, ProviderInfo>,
  ): CalendarData => {
    const data: Record<string, Activity[]> = Object.fromEntries(
      Object.keys(providers).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).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).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;
  };