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

const calculateLevel = (count: number, avgCount: number): number => {
  if (count === 0) return 0;
  if (count <= avgCount * 0.5) return 1;
  if (count <= avgCount) return 2;
  if (count <= avgCount * 1.5) return 3;
  return 4;
};

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.setDate(startDate.getDate() - 364);
  const dayOfWeek = startDate.getDay();
  startDate.setDate(startDate.getDate() - dayOfWeek);

  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;
      }
    });
  });

  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 });
    });
  }

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

  Object.entries(data).forEach(([provider, days]) => {
    const avgCount = avgCounts[provider];
    for (const day of days) {
      day.level = calculateLevel(day.count, avgCount);
    }
  });

  return data;
};

export const compactCalendarData = (data: CalendarData): CompactCalendarData => {
  const compact: CompactCalendarData = {};
  for (const [key, activities] of Object.entries(data)) {
    if (activities.length > 0) {
      compact[key] = {
        startDate: activities[0].date,
        counts: activities.map((a) => a.count),
      };
    }
  }
  return compact;
};

export const expandCalendarData = (compact: CompactCalendarData): CalendarData => {
  const data: CalendarData = {};
  for (const [key, { startDate, counts }] of Object.entries(compact)) {
    const total = counts.reduce((sum, c) => sum + c, 0);
    const avg = total / counts.length || 0;
    const start = new Date(startDate + "T00:00:00");

    const activities: Activity[] = counts.map((count, i) => {
      const date = new Date(start);
      date.setDate(date.getDate() + i);
      return {
        date: date.toISOString().split("T")[0],
        count,
        level: calculateLevel(count, avg),
      };
    });

    data[key] = activities;
  }
  return data;
};