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

'use client';

import {
  BarChart,
  Bar,
  XAxis,
  YAxis,
  CartesianGrid,
  Tooltip,
  ResponsiveContainer,
  PieChart,
  Pie,
  Cell,
  Legend,
} from 'recharts';
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 = [
  'hsl(var(--chart-1))',
  'hsl(var(--chart-2))',
  'hsl(var(--chart-3))',
  'hsl(var(--chart-4))',
  'hsl(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>
        <ResponsiveContainer width="100%" height={250}>
          <BarChart data={data} layout="vertical" margin={{ left: 20, right: 20 }}>
            <CartesianGrid strokeDasharray="3 3" stroke="hsl(var(--border))" />
            <XAxis type="number" stroke="hsl(var(--muted-foreground))" fontSize={12} />
            <YAxis
              type="category"
              dataKey="name"
              stroke="hsl(var(--muted-foreground))"
              fontSize={11}
              width={120}
            />
            <Tooltip
              contentStyle={{
                backgroundColor: 'hsl(var(--card))',
                border: '1px solid hsl(var(--border))',
                borderRadius: 8,
                fontSize: 12,
              }}
            />
            <Bar dataKey="count" fill="hsl(var(--chart-1))" radius={[0, 4, 4, 0]} />
          </BarChart>
        </ResponsiveContainer>
      </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>
        <ResponsiveContainer width="100%" height={250}>
          <PieChart>
            <Pie
              data={data}
              cx="50%"
              cy="50%"
              innerRadius={50}
              outerRadius={90}
              paddingAngle={2}
              dataKey="value"
              label={({ name, percent }) => `${name} (${((percent ?? 0) * 100).toFixed(0)}%)`}
              labelLine={false}
              fontSize={11}
            >
              {data.map((_, idx) => (
                <Cell key={idx} fill={COLORS[idx % COLORS.length]} />
              ))}
            </Pie>
            <Tooltip
              contentStyle={{
                backgroundColor: 'hsl(var(--card))',
                border: '1px solid hsl(var(--border))',
                borderRadius: 8,
                fontSize: 12,
              }}
            />
            <Legend wrapperStyle={{ fontSize: 11, color: 'hsl(var(--muted-foreground))' }} />
          </PieChart>
        </ResponsiveContainer>
      </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>
  );
}