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

export const generateCalendarData = (
  modelData: ModelData[],
  providers: ProviderInfo[]
): CalendarData => {
  const data: Record<string, Activity[]> = Object.fromEntries(
    providers.map((provider) => [provider.authors[0], []]),
  );

  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];
    providers.forEach(({ authors }) => {
      if (authors.some((author) => item.id.startsWith(author))) {
        countMap[authors[0]] = countMap[authors[0]] || {};
        countMap[authors[0]][dateString] =
          (countMap[authors[0]][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];

    providers.forEach(({ authors }) => {
      const count = countMap[authors[0]]?.[dateString] || 0;
      data[authors[0]].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;
};