#!/usr/bin/env python3

import argparse
import json
import math
import sqlite3
import sys
from collections import Counter, defaultdict
from pathlib import Path

import numpy as np
from sklearn.cluster import HDBSCAN


EMBEDDING_DIMENSIONS = 128
EMBEDDING_BYTES = EMBEDDING_DIMENSIONS * 4


def parse_arguments():
    parser = argparse.ArgumentParser(
        description=(
            "Esegue HDBSCAN sugli embedding facciali "
            "senza modificare il database."
        ),
    )

    parser.add_argument(
        "--database",
        required=True,
        help="Percorso del database SQLite",
    )

    parser.add_argument(
        "--user",
        required=True,
        help="Utente proprietario dei volti",
    )

    parser.add_argument(
        "--min-cluster-size",
        type=int,
        default=2,
    )

    parser.add_argument(
        "--min-samples",
        type=int,
        default=1,
    )

    parser.add_argument(
        "--metric",
        default="euclidean",
    )

    parser.add_argument(
        "--selection-method",
        choices=("eom", "leaf"),
        default="eom",
    )

    parser.add_argument(
        "--scope",
        choices=("all", "unclustered"),
        default="all",
    )

    parser.add_argument(
        "--pretty",
        action="store_true",
    )

    arguments = parser.parse_args()

    if arguments.min_cluster_size < 2:
        parser.error(
            "--min-cluster-size deve essere almeno 2"
        )

    if arguments.min_samples < 1:
        parser.error(
            "--min-samples deve essere almeno 1"
        )

    return arguments


def load_faces(database_path, scope, user):
    connection = sqlite3.connect(database_path)
    connection.row_factory = sqlite3.Row

    query = """
        SELECT
            f.id,
            f.photo_id,
            f.face_index,
            f.score,
            f.embedding,
            p.name AS photo_name,
            p.path AS photo_path,
            COALESCE(p.rotation, 0) AS photo_rotation
        FROM faces AS f
        JOIN photos AS p
          ON p.id = f.photo_id
        WHERE f.embedding IS NOT NULL
          AND p.user = ?
    """

    parameters = [user]

    if scope == "unclustered":
        query += """
          AND f.cluster_id IS NULL
        """

    query += """
        ORDER BY
            f.id
    """

    rows = connection.execute(
        query,
        parameters,
    ).fetchall()

    connection.close()

    if len(rows) < 2:
        raise RuntimeError(
            "Sono necessari almeno due embedding"
        )

    faces = []
    embeddings = []

    for row in rows:
        blob = row["embedding"]

        if len(blob) != EMBEDDING_BYTES:
            raise RuntimeError(
                "Embedding non valido per "
                f"{row['id']}: {len(blob)} byte"
            )

        vector = np.frombuffer(
            blob,
            dtype="<f4",
            count=EMBEDDING_DIMENSIONS,
        ).astype(np.float64)

        if not np.all(np.isfinite(vector)):
            raise RuntimeError(
                f"Embedding non finito: {row['id']}"
            )

        norm = float(np.linalg.norm(vector))

        if not math.isfinite(norm) or norm <= 0.0:
            raise RuntimeError(
                f"Norma non valida: {row['id']}"
            )

        vector /= norm

        faces.append(
            {
                "id": row["id"],
                "photo_id": row["photo_id"],
                "face_index": row["face_index"],
                "detector_score": row["score"],
                "photo_name": row["photo_name"],
                "photo_path": row["photo_path"],
                "photo_rotation": row["photo_rotation"],
                "input_norm": norm,
            }
        )

        embeddings.append(vector)

    matrix = np.vstack(embeddings)

    return faces, matrix


def cosine_statistics(indices, matrix):
    if len(indices) < 2:
        return {
            "pair_count": 0,
            "minimum": None,
            "maximum": None,
            "mean": None,
        }

    values = []

    for left_position in range(len(indices)):
        for right_position in range(
            left_position + 1,
            len(indices),
        ):
            left_index = indices[left_position]
            right_index = indices[right_position]

            similarity = float(
                np.dot(
                    matrix[left_index],
                    matrix[right_index],
                )
            )

            values.append(similarity)

    return {
        "pair_count": len(values),
        "minimum": float(min(values)),
        "maximum": float(max(values)),
        "mean": float(sum(values) / len(values)),
    }


def build_result(arguments, faces, matrix, labels, probabilities):
    label_counts = Counter(
        int(label) for label in labels
    )

    cluster_indices = defaultdict(list)

    for index, label_value in enumerate(labels):
        label = int(label_value)

        if label >= 0:
            cluster_indices[label].append(index)

    clusters = []

    for label in sorted(cluster_indices):
        indices = cluster_indices[label]

        members = []

        for index in indices:
            face = dict(faces[index])

            face["probability"] = float(
                probabilities[index]
            )

            members.append(face)

        clusters.append(
            {
                "label": label,
                "size": len(indices),
                "probability": {
                    "minimum": float(
                        min(
                            probabilities[index]
                            for index in indices
                        )
                    ),
                    "maximum": float(
                        max(
                            probabilities[index]
                            for index in indices
                        )
                    ),
                    "mean": float(
                        sum(
                            probabilities[index]
                            for index in indices
                        )
                        / len(indices)
                    ),
                },
                "cosine_similarity": cosine_statistics(
                    indices,
                    matrix,
                ),
                "members": members,
            }
        )

    noise_members = []

    for index, label_value in enumerate(labels):
        if int(label_value) != -1:
            continue

        face = dict(faces[index])
        face["probability"] = float(
            probabilities[index]
        )

        noise_members.append(face)

    total_faces = len(faces)
    noise_count = label_counts.get(-1, 0)

    return {
        "dry_run": True,
        "database": str(arguments.database),
        "parameters": {
            "min_cluster_size":
                arguments.min_cluster_size,
            "min_samples":
                arguments.min_samples,
            "metric": arguments.metric,
            "cluster_selection_method":
                arguments.selection_method,
            "scope": arguments.scope,
            "user": arguments.user,
        },
        "summary": {
            "total_faces": total_faces,
            "cluster_count": len(clusters),
            "clustered_faces":
                total_faces - noise_count,
            "noise_faces": noise_count,
            "noise_percent": (
                100.0 * noise_count / total_faces
            ),
        },
        "clusters": clusters,
        "noise": noise_members,
    }


def main():
    arguments = parse_arguments()

    database_path = Path(
        arguments.database
    ).expanduser().resolve()

    if not database_path.is_file():
        raise RuntimeError(
            f"Database non trovato: {database_path}"
        )

    faces, matrix = load_faces(
        database_path,
        arguments.scope,
        arguments.user,
    )

    clusterer = HDBSCAN(
        min_cluster_size=
            arguments.min_cluster_size,
        min_samples=arguments.min_samples,
        metric=arguments.metric,
        cluster_selection_method=
            arguments.selection_method,
        allow_single_cluster=False,
        copy=True,
    )

    labels = clusterer.fit_predict(matrix)
    probabilities = clusterer.probabilities_

    result = build_result(
        arguments,
        faces,
        matrix,
        labels,
        probabilities,
    )

    json.dump(
        result,
        sys.stdout,
        indent=2 if arguments.pretty else None,
        ensure_ascii=False,
    )

    sys.stdout.write("\n")


if __name__ == "__main__":
    try:
        main()
    except Exception as error:
        print(
            f"[face_hdbscan] errore: {error}",
            file=sys.stderr,
        )

        raise