import datasets
import evaluate
import numpy as np
from PIL import Image


_DESCRIPTION = """
Peak Signal-to-Noise Ratio (PSNR) is a metric used to quantify the quality of a reconstructed or compressed image compared to its original version. It measures the ratio between the maximum possible power of a signal (the original image) and the power of corrupting noise (the error introduced by compression or reconstruction). PSNR is expressed in decibels (dB), with higher values indicating better quality.

PSNR is particularly useful in image and video compression, reconstruction, and quality assessment tasks. It provides an objective measure of the fidelity of the compressed image with respect to the original image. The PSNR can be calculated using the formula:

    PSNR = 20 * log10(MAX_I / sqrt(MSE))

Where:
- MAX_I is the maximum possible pixel value of the image (for example, 255 for an 8-bit image).
- MSE is the Mean Squared Error between the original and the reconstructed image, defined as:
    MSE = (1 / (m * n)) * sum((I1(i,j) - I2(i,j))^2)

Here:
- I1 is the original image,
- I2 is the reconstructed image,
- m and n represent the image dimensions (height and width),
- and the summation runs over all pixel locations (i, j).

Characteristics of PSNR:
Higher PSNR values indicate higher image quality, with infinity representing identical images.
Typical PSNR values range between 20 and 50 dB, depending on the level of distortion.
Applications: Image and video compression (e.g., JPEG, MPEG), reconstruction (e.g., super-resolution), and quality assessment.
"""

_KWARGS_DESCRIPTION = """
Args:
    predictions (`list` of `PIL.Image` or `numpy.ndarray`): Predicted images.
    references (`list` of `PIL.Image` or `numpy.ndarray`): Ground truth images.
Returns:
    psnr (`float`): PSNR score. Higher values indicate better quality, with infinity for identical images. Typical values range between 20 and 50 dB.
Examples:
    Example 
        >>> psnr_metric = evaluate.load("Bekhouche/PSNR")
        >>> results = psnr_metric.compute(
                references=[Image.open('reference1.png'), Image.open('reference2.png')], 
                predictions=[Image.open('prediction1.png'), Image.open('prediction2.png')]
            )
        >>> print(results)
        {'psnr': 30.5}
"""

_CITATION = """
"""

@evaluate.utils.file_utils.add_start_docstrings(_DESCRIPTION, _KWARGS_DESCRIPTION)
class PSNR(evaluate.Metric):
    def _info(self):
        return evaluate.MetricInfo(
            module_type="metric",
            description=_DESCRIPTION,
            citation=_CITATION,
            inputs_description=_KWARGS_DESCRIPTION,
            features=datasets.Features(
                {
                    "predictions": datasets.Sequence(datasets.Image(decode=True)),
                    "references": datasets.Sequence(datasets.Image(decode=True)),
                }
                if self.config_name == "multilabel"
                else {
                    "predictions": datasets.Image(decode=True),
                    "references": datasets.Image(decode=True)
                }
            ),
            reference_urls=["https://huggingface.co/spaces/Bekhouche/PSNR"],
        )

    def _compute(self, predictions, references):
        if isinstance(predictions, list) and isinstance(references, list):
            psnr = sum([self._compute_psnr(prediction, reference) for prediction, reference in zip(predictions, references)]) / len(predictions)
        elif (isinstance(predictions, (Image.Image, np.ndarray)) and 
              isinstance(references, (Image.Image, np.ndarray))):
            psnr = self._compute_psnr(predictions, references)
        else:
            raise ValueError("Predictions and references must be either lists of images (PIL or numpy arrays) or single images.")
        return {
            "psnr": float(psnr)
        }
    
    @staticmethod
    def _compute_psnr(prediction, reference):
        if isinstance(prediction, Image.Image):
            prediction = np.array(prediction)
        if isinstance(reference, Image.Image):
            reference = np.array(reference)

        if prediction.shape != reference.shape:
            raise ValueError("Images must have the same dimensions and number of channels for PSNR calculation.")
        
        mse = np.mean((prediction - reference) ** 2)
        if mse == 0:
            psnr = float('inf')
        else:
            max_pixel_value = 255.0
            psnr = 20 * np.log10(max_pixel_value / np.sqrt(mse))
        return psnr