learning_ai_common_plat/dashboards/admin-web/src/components/extraction/entity-chart.tsx

'use client';

import { BarChart, Donut } from '@/components/charts';
import { Card, CardContent, CardHeader, CardTitle } from '@/components/ui/card';

interface ExtractionEntity {
  extraction_class: string;
  extraction_text: string;
  attributes?: Record<string, string>;
}

interface EntityChartProps {
  extractions: ExtractionEntity[];
  title?: string;
}

const COLORS = [
  'var(--chart-1)',
  'var(--chart-2)',
  'var(--chart-3)',
  'var(--chart-4)',
  'var(--chart-5)',
];

export function EntityBarChart({ extractions, title = 'Entities by Class' }: EntityChartProps) {
  const classCounts = extractions.reduce(
    (acc, e) => {
      acc[e.extraction_class] = (acc[e.extraction_class] || 0) + 1;
      return acc;
    },
    {} as Record<string, number>
  );

  const data = Object.entries(classCounts)
    .map(([name, count]) => ({ name, count }))
    .sort((a, b) => b.count - a.count);

  if (data.length === 0) return null;

  return (
    <Card>
      <CardHeader>
        <CardTitle className="text-sm">{title}</CardTitle>
      </CardHeader>
      <CardContent>
        <BarChart
          data={data.map(d => ({ id: d.name, value: d.count, label: d.name }))}
          width={640}
          height={250}
          className="h-auto w-full"
          ariaLabel={title}
        />
      </CardContent>
    </Card>
  );
}

export function EntityPieChart({ extractions, title = 'Class Distribution' }: EntityChartProps) {
  const classCounts = extractions.reduce(
    (acc, e) => {
      acc[e.extraction_class] = (acc[e.extraction_class] || 0) + 1;
      return acc;
    },
    {} as Record<string, number>
  );

  const data = Object.entries(classCounts)
    .map(([name, value]) => ({ name, value }))
    .sort((a, b) => b.value - a.value);

  if (data.length === 0) return null;

  return (
    <Card>
      <CardHeader>
        <CardTitle className="text-sm">{title}</CardTitle>
      </CardHeader>
      <CardContent className="flex justify-center">
        <Donut
          slices={data.map((d, idx) => ({
            id: d.name,
            label: d.name,
            value: d.value,
            color: COLORS[idx % COLORS.length],
          }))}
          size={220}
          ariaLabel={title}
        />
      </CardContent>
    </Card>
  );
}

export function EntityTimeline({ extractions }: { extractions: ExtractionEntity[] }) {
  if (extractions.length === 0) return null;

  return (
    <Card>
      <CardHeader>
        <CardTitle className="text-sm">Entity Timeline</CardTitle>
      </CardHeader>
      <CardContent>
        <div className="relative pl-6 space-y-3">
          <div className="absolute left-2 top-0 bottom-0 w-px bg-border" />
          {extractions.slice(0, 20).map((e, i) => (
            <div key={i} className="relative flex items-start gap-3">
              <div
                className="absolute -left-4 mt-1.5 h-2.5 w-2.5 rounded-full"
                style={{
                  backgroundColor:
                    COLORS[
                      Object.keys(
                        extractions.reduce(
                          (acc, ex) => {
                            acc[ex.extraction_class] = true;
                            return acc;
                          },
                          {} as Record<string, boolean>
                        )
                      ).indexOf(e.extraction_class) % COLORS.length
                    ],
                }}
              />
              <div className="min-w-0 flex-1">
                <span className="text-[10px] font-medium uppercase tracking-wider text-muted-foreground">
                  {e.extraction_class}
                </span>
                <p className="text-sm text-foreground">{e.extraction_text}</p>
              </div>
            </div>
          ))}
        </div>
      </CardContent>
    </Card>
  );
}