import warnings

import numpy as np
import pandas as pd
from matplotlib import pyplot as plt
from sklearn.metrics import pairwise_distances

warnings.filterwarnings('''ignore''')

UpperCAmelCase							= '''K-MEANS-CLUST/ '''


def SCREAMING_SNAKE_CASE (   lowercase_:		Tuple						,  lowercase_:		Optional[int]						,  lowercase_:		str=None ):
       if seed is not None:  # useful for obtaining consistent results
              np.random.seed(lowercase_ )
       _UpperCamelCase   :				str       =     data.shape[0]  # number of data points

       # Pick K indices from range [0, N).
       _UpperCamelCase   :				Optional[int]       =     np.random.randint(0						,  lowercase_						,  lowercase_ )

       # Keep centroids as dense format, as many entries will be nonzero due to averaging.
       # As long as at least one document in a cluster contains a word,
       # it will carry a nonzero weight in the TF-IDF vector of the centroid.
       _UpperCamelCase   :				int       =     data[rand_indices, :]

       return centroids


def SCREAMING_SNAKE_CASE (   lowercase_:		Optional[Any]						,  lowercase_:		List[Any] ):
       return pairwise_distances(lowercase_						,  lowercase_						,  metric='euclidean' )


def SCREAMING_SNAKE_CASE (   lowercase_:		List[str]						,  lowercase_:		Optional[int] ):
       # Compute distances between each data point and the set of centroids:
       # Fill in the blank (RHS only)
       _UpperCamelCase   :				Any       =     centroid_pairwise_dist(lowercase_						,  lowercase_ )

       # Compute cluster assignments for each data point:
       # Fill in the blank (RHS only)
       _UpperCamelCase   :				List[str]       =     np.argmin(lowercase_						,  axis=1 )

       return cluster_assignment


def SCREAMING_SNAKE_CASE (   lowercase_:		Optional[int]						,  lowercase_:		Union[str, Any]						,  lowercase_:		List[Any] ):
       _UpperCamelCase   :				List[Any]       =     []
       for i in range(lowercase_ ):
              # Select all data points that belong to cluster i. Fill in the blank (RHS only)
              _UpperCamelCase   :				Union[str, Any]       =     data[cluster_assignment == i]
              # Compute the mean of the data points. Fill in the blank (RHS only)
              _UpperCamelCase   :				str       =     member_data_points.mean(axis=0 )
              new_centroids.append(lowercase_ )
       _UpperCamelCase   :				Any       =     np.array(lowercase_ )

       return new_centroids


def SCREAMING_SNAKE_CASE (   lowercase_:		Optional[Any]						,  lowercase_:		List[str]						,  lowercase_:		Dict						,  lowercase_:		int ):
       _UpperCamelCase   :				Dict       =     0.0
       for i in range(lowercase_ ):
              # Select all data points that belong to cluster i. Fill in the blank (RHS only)
              _UpperCamelCase   :				List[str]       =     data[cluster_assignment == i, :]

              if member_data_points.shape[0] > 0:  # check if i-th cluster is non-empty
                     # Compute distances from centroid to data points (RHS only)
                     _UpperCamelCase   :				Union[str, Any]       =     pairwise_distances(
                         lowercase_						,  [centroids[i]]						,  metric='euclidean' )
                     _UpperCamelCase   :				Tuple       =     distances**2
                     heterogeneity += np.sum(lowercase_ )

       return heterogeneity


def SCREAMING_SNAKE_CASE (   lowercase_:		Optional[int]						,  lowercase_:		int ):
       plt.figure(figsize=(7, 4) )
       plt.plot(lowercase_						,  linewidth=4 )
       plt.xlabel('# Iterations' )
       plt.ylabel('Heterogeneity' )
       plt.title(F"Heterogeneity of clustering over time, K={k:d}" )
       plt.rcParams.update({'font.size': 16} )
       plt.show()


def SCREAMING_SNAKE_CASE (   lowercase_:		str						,  lowercase_:		str						,  lowercase_:		List[Any]						,  lowercase_:		Optional[int]=500						,  lowercase_:		Dict=None						,  lowercase_:		Any=False ):
       _UpperCamelCase   :				Optional[int]       =     initial_centroids[:]
       _UpperCamelCase   :				str       =     None

       for itr in range(lowercase_ ):
              if verbose:
                     print(lowercase_						,  end='' )

              # 1. Make cluster assignments using nearest centroids
              _UpperCamelCase   :				Optional[Any]       =     assign_clusters(lowercase_						,  lowercase_ )

              # 2. Compute a new centroid for each of the k clusters, averaging all data
              #    points assigned to that cluster.
              _UpperCamelCase   :				List[Any]       =     revise_centroids(lowercase_						,  lowercase_						,  lowercase_ )

              # Check for convergence: if none of the assignments changed, stop
              if (
                  prev_cluster_assignment is not None
                  and (prev_cluster_assignment == cluster_assignment).all()
              ):
                     break

              # Print number of new assignments
              if prev_cluster_assignment is not None:
                     _UpperCamelCase   :				Any       =     np.sum(prev_cluster_assignment != cluster_assignment )
                     if verbose:
                            print(
                                F"    {num_changed:5d} elements changed their cluster assignment." )

        # Record heterogeneity convergence metric
              if record_heterogeneity is not None:
                     # YOUR CODE HERE
                     _UpperCamelCase   :				str       =     compute_heterogeneity(lowercase_						,  lowercase_						,  lowercase_						,  lowercase_ )
                     record_heterogeneity.append(lowercase_ )

              _UpperCamelCase   :				Optional[int]       =     cluster_assignment[:]

       return centroids, cluster_assignment


# Mock test below
if False:  # change to true to run this test case.
 from sklearn import datasets as ds

 UpperCAmelCase							= ds.load_iris()
 UpperCAmelCase							= 3
 UpperCAmelCase							= []
 UpperCAmelCase							= get_initial_centroids(dataset['''data'''], k, seed=0)
 UpperCAmelCase						,    UpperCAmelCase							= kmeans(
     dataset['''data'''],
     k,
     initial_centroids,
     maxiter=4_00,
     record_heterogeneity=heterogeneity,
     verbose=True,
 )
 plot_heterogeneity(heterogeneity, k)


def SCREAMING_SNAKE_CASE (   lowercase_:		pd.DataFrame						,  lowercase_:		np.ndarray						,  lowercase_:		Any=None ):
       # Fill missing values with given rules
       if fill_missing_report:
              _UpperCamelCase   :				int       =     df.fillna(value=lowercase_ )
       _UpperCamelCase   :				Dict       =     1
       _UpperCamelCase   :				List[Any]       =     df.select_dtypes(np.number ).columns
       _UpperCamelCase   :				str       =     (
           df.groupby(['Cluster'] )[  # construct report dataframe
               numeric_cols
           ]  # group by cluster number
           .agg(
               [
                   ('sum', np.sum),
                   ('mean_with_zeros', lambda lowercase_ : np.mean(np.nan_to_num(lowercase_ ) )),
                   ('mean_without_zeros', lambda lowercase_ : x.replace(0						,  np.NaN ).mean()),
                   (
                       'mean_25-75',
                       lambda lowercase_ : np.mean(
                           np.nan_to_num(
                               sorted(lowercase_ )[
                                   round(len(lowercase_ ) * 25 / 100 ) : round(len(lowercase_ ) * 75 / 100 )
                               ] ) ),
                   ),
                   ('mean_with_na', np.mean),
                   ('min', lambda lowercase_ : x.min()),
                   ('5%', lambda lowercase_ : x.quantile(0.05 )),
                   ('25%', lambda lowercase_ : x.quantile(0.25 )),
                   ('50%', lambda lowercase_ : x.quantile(0.50 )),
                   ('75%', lambda lowercase_ : x.quantile(0.75 )),
                   ('95%', lambda lowercase_ : x.quantile(0.95 )),
                   ('max', lambda lowercase_ : x.max()),
                   ('count', lambda lowercase_ : x.count()),
                   ('stdev', lambda lowercase_ : x.std()),
                   ('mode', lambda lowercase_ : x.mode()[0]),
                   ('median', lambda lowercase_ : x.median()),
                   ('# > 0', lambda lowercase_ : (x > 0).sum()),
               ] )
           .T.reset_index()
           .rename(index=lowercase_						,  columns={'level_0': 'Features', 'level_1': 'Type'} )
       )  # rename columns
       # calculate the size of cluster(count of clientID's)
       _UpperCamelCase   :				int       =     report[
           (report['Features'] == 'dummy') & (report['Type'] == 'count')
       ].copy()  # avoid SettingWithCopyWarning
       _UpperCamelCase   :				int       =     (
           'ClusterSize'  # rename created cluster df to match report column names
       )
       _UpperCamelCase   :				List[Any]       =     '# of Customers'
       _UpperCamelCase   :				int       =     pd.DataFrame(
           clustersize.iloc[:, 2:].values
           / clustersize.iloc[:, 2:].values.sum() )
       _UpperCamelCase   :				Optional[int]       =     '% of Customers'  # rename created cluster df to match report column names
       _UpperCamelCase   :				List[str]       =     'ClusterProportion'
       _UpperCamelCase   :				str       =     clusterproportion.columns.tolist()
       _UpperCamelCase   :				Any       =     cols[-2:] + cols[:-2]
       _UpperCamelCase   :				List[Any]       =     clusterproportion[cols]  # rearrange columns to match report
       _UpperCamelCase   :				List[str]       =     report.columns
       _UpperCamelCase   :				int       =     pd.DataFrame(
           abs(
               report[report['Type'] == 'count'].iloc[:, 2:].values
               - clustersize.iloc[:, 2:].values ) )  # generating df with count of nan values
       _UpperCamelCase   :				List[str]       =     0
       _UpperCamelCase   :				Optional[Any]       =     '# of nan'
       _UpperCamelCase   :				int       =     report[
           report['Type'] == 'count'
       ].Features.tolist()  # filling values in order to match report
       _UpperCamelCase   :				int       =     a.columns.tolist()
       _UpperCamelCase   :				int       =     cols[-2:] + cols[:-2]
       _UpperCamelCase   :				Optional[Any]       =     a[cols]  # rearrange columns to match report
       _UpperCamelCase   :				int       =     report.columns  # rename columns to match report
       _UpperCamelCase   :				Any       =     report.drop(
           report[report.Type == 'count'].index )  # drop count values except cluster size
       _UpperCamelCase   :				Optional[Any]       =     pd.concat(
           [report, a, clustersize, clusterproportion]						,  axis=0 )  # concat report with clustert size and nan values
       _UpperCamelCase   :				Tuple       =     report['Features'].isin(lowercase_ )
       _UpperCamelCase   :				Optional[int]       =     report.columns.tolist()
       _UpperCamelCase   :				str       =     cols[0:2] + cols[-1:] + cols[2:-1]
       _UpperCamelCase   :				Union[str, Any]       =     report[cols]
       _UpperCamelCase   :				List[str]       =     {
           'ClusterSize': 9,
           'ClusterProportion': 8,
           'mean_with_zeros': 7,
           'mean_with_na': 6,
           'max': 5,
           '50%': 4,
           'min': 3,
           '25%': 2,
           '75%': 1,
           '# of nan': 0,
           '# > 0': -1,
           'sum_with_na': -2,
       }
       _UpperCamelCase   :				List[Any]       =     (
           report.assign(
               Sortera=lambda lowercase_ : x.Type.map(lowercase_ )						,  Sortera=lambda lowercase_ : list(reversed(range(len(lowercase_ ) ) ) )						,  )
           .sort_values(['Sorter1', 'Mark', 'Sorter2']						,  ascending=lowercase_ )
           .drop(['Sorter1', 'Sorter2']						,  axis=1 )
       )
       _UpperCamelCase   :				Any       =     ''
       _UpperCamelCase   :				int       =     report.reset_index()
       _UpperCamelCase   :				Optional[Any]       =     report.drop(columns=['index'] )
       return report


if __name__ == "__main__":
 import doctest

 doctest.testmod()




import gc
import unittest

import numpy as np
import torch
from transformers import CLIPTextConfig, CLIPTextModel, CLIPTokenizer

from diffusers import (
    AutoencoderKL,
    DDIMParallelScheduler,
    DDPMParallelScheduler,
    StableDiffusionParadigmsPipeline,
    UNetaDConditionModel,
)
from diffusers.utils import slow, torch_device
from diffusers.utils.testing_utils import (
    enable_full_determinism,
    require_torch_gpu,
)

from ..pipeline_params import TEXT_TO_IMAGE_BATCH_PARAMS, TEXT_TO_IMAGE_IMAGE_PARAMS, TEXT_TO_IMAGE_PARAMS
from ..test_pipelines_common import PipelineLatentTesterMixin, PipelineTesterMixin


enable_full_determinism()




class       __lowerCAmelCase     (				a_,		a_,		unittest.TestCase  ):



 '''simple docstring'''
 _lowerCamelCase  :		Any 			=       StableDiffusionParadigmsPipeline
 _lowerCamelCase  :		Any 			=       TEXT_TO_IMAGE_PARAMS
 _lowerCamelCase  :		str 			=       TEXT_TO_IMAGE_BATCH_PARAMS
 _lowerCamelCase  :		int 			=       TEXT_TO_IMAGE_IMAGE_PARAMS
 _lowerCamelCase  :		List[Any] 			=       TEXT_TO_IMAGE_IMAGE_PARAMS

 def     __snake_case    (	self				:						int     )       ->							List[str]:
        torch.manual_seed(0     )
        _UpperCamelCase   :				Union[str, Any]       =     UNetaDConditionModel(
            block_out_channels=(3_2, 6_4)				,	layers_per_block=2				,	sample_size=3_2				,	in_channels=4				,	out_channels=4				,	down_block_types=('DownBlock2D', 'CrossAttnDownBlock2D')				,	up_block_types=('CrossAttnUpBlock2D', 'UpBlock2D')				,	cross_attention_dim=3_2				,	attention_head_dim=(2, 4)				,	use_linear_projection=lowerCamelCase				,	)
        _UpperCamelCase   :				str       =     DDIMParallelScheduler(
            beta_start=0.0_0_0_8_5				,	beta_end=0.0_1_2				,	beta_schedule='scaled_linear'				,	clip_sample=lowerCamelCase				,	set_alpha_to_one=lowerCamelCase				,	)
        torch.manual_seed(0     )
        _UpperCamelCase   :				str       =     AutoencoderKL(
            block_out_channels=[3_2, 6_4]				,	in_channels=3				,	out_channels=3				,	down_block_types=['DownEncoderBlock2D', 'DownEncoderBlock2D']				,	up_block_types=['UpDecoderBlock2D', 'UpDecoderBlock2D']				,	latent_channels=4				,	sample_size=1_2_8				,	)
        torch.manual_seed(0     )
        _UpperCamelCase   :				int       =     CLIPTextConfig(
            bos_token_id=0				,	eos_token_id=2				,	hidden_size=3_2				,	intermediate_size=3_7				,	layer_norm_eps=1e-05				,	num_attention_heads=4				,	num_hidden_layers=5				,	pad_token_id=1				,	vocab_size=1_0_0_0				,	hidden_act='gelu'				,	projection_dim=5_1_2				,	)
        _UpperCamelCase   :				Any       =     CLIPTextModel(lowerCamelCase     )
        _UpperCamelCase   :				Optional[Any]       =     CLIPTokenizer.from_pretrained('hf-internal-testing/tiny-random-clip'     )

        _UpperCamelCase   :				Tuple       =     {
            'unet': unet,
            'scheduler': scheduler,
            'vae': vae,
            'text_encoder': text_encoder,
            'tokenizer': tokenizer,
            'safety_checker': None,
            'feature_extractor': None,
        }
        return components

 def     __snake_case    (	self				:						List[Any]				,	lowerCamelCase				:						Any				,	lowerCamelCase				:						str=0     )       ->							Dict:
        if str(lowerCamelCase     ).startswith('mps'     ):
               _UpperCamelCase   :				List[Any]       =     torch.manual_seed(lowerCamelCase     )
        else:
               _UpperCamelCase   :				Union[str, Any]       =     torch.Generator(device=lowerCamelCase     ).manual_seed(lowerCamelCase     )
        _UpperCamelCase   :				Optional[Any]       =     {
            'prompt': 'a photograph of an astronaut riding a horse',
            'generator': generator,
            'num_inference_steps': 1_0,
            'guidance_scale': 6.0,
            'output_type': 'numpy',
            'parallel': 3,
            'debug': True,
        }
        return inputs

 def     __snake_case    (	self				:						List[Any]     )       ->							str:
        _UpperCamelCase   :				List[str]       =     'cpu'  # ensure determinism for the device-dependent torch.Generator
        _UpperCamelCase   :				List[Any]       =     self.get_dummy_components()
        _UpperCamelCase   :				Tuple       =     StableDiffusionParadigmsPipeline(**lowerCamelCase     )
        _UpperCamelCase   :				str       =     sd_pipe.to(lowerCamelCase     )
        sd_pipe.set_progress_bar_config(disable=lowerCamelCase     )

        _UpperCamelCase   :				Any       =     self.get_dummy_inputs(lowerCamelCase     )
        _UpperCamelCase   :				str       =     sd_pipe(**lowerCamelCase     ).images
        _UpperCamelCase   :				str       =     image[0, -3:, -3:, -1]
        assert image.shape == (1, 6_4, 6_4, 3)

        _UpperCamelCase   :				Any       =     np.array([0.4_7_7_3, 0.5_4_1_7, 0.4_7_2_3, 0.4_9_2_5, 0.5_6_3_1, 0.4_7_5_2, 0.5_2_4_0, 0.4_9_3_5, 0.5_0_2_3]     )

        assert np.abs(image_slice.flatten() - expected_slice     ).max() < 1e-2

 def     __snake_case    (	self				:						List[Any]     )       ->							Union[str, Any]:
        _UpperCamelCase   :				Dict       =     'cpu'  # ensure determinism for the device-dependent torch.Generator
        _UpperCamelCase   :				str       =     self.get_dummy_components()
        torch.manual_seed(0     )
        _UpperCamelCase   :				Union[str, Any]       =     DDPMParallelScheduler()
        torch.manual_seed(0     )
        _UpperCamelCase   :				List[Any]       =     StableDiffusionParadigmsPipeline(**lowerCamelCase     )
        _UpperCamelCase   :				str       =     sd_pipe.to(lowerCamelCase     )
        sd_pipe.set_progress_bar_config(disable=lowerCamelCase     )

        _UpperCamelCase   :				int       =     self.get_dummy_inputs(lowerCamelCase     )
        _UpperCamelCase   :				Optional[int]       =     sd_pipe(**lowerCamelCase     ).images
        _UpperCamelCase   :				Dict       =     image[0, -3:, -3:, -1]
        assert image.shape == (1, 6_4, 6_4, 3)

        _UpperCamelCase   :				List[Any]       =     np.array([0.3_5_7_3, 0.4_4_2_0, 0.4_9_6_0, 0.4_7_9_9, 0.3_7_9_6, 0.3_8_7_9, 0.4_8_1_9, 0.4_3_6_5, 0.4_4_6_8]     )

        assert np.abs(image_slice.flatten() - expected_slice     ).max() < 1e-2

 def     __snake_case    (	self				:						Any     )       ->							List[Any]:
        super().test_inference_batch_consistent(batch_sizes=[1, 2]     )

 def     __snake_case    (	self				:						Union[str, Any]     )       ->							Tuple:
        super().test_inference_batch_single_identical(batch_size=2				,	expected_max_diff=3e-3     )

 def     __snake_case    (	self				:						List[str]     )       ->							List[str]:
        _UpperCamelCase   :				List[str]       =     'cpu'  # ensure determinism for the device-dependent torch.Generator
        _UpperCamelCase   :				Tuple       =     self.get_dummy_components()
        _UpperCamelCase   :				Any       =     StableDiffusionParadigmsPipeline(**lowerCamelCase     )
        _UpperCamelCase   :				List[Any]       =     sd_pipe.to(lowerCamelCase     )
        sd_pipe.set_progress_bar_config(disable=lowerCamelCase     )

        _UpperCamelCase   :				Tuple       =     self.get_dummy_inputs(lowerCamelCase     )
        _UpperCamelCase   :				int       =     'french fries'
        _UpperCamelCase   :				Optional[Any]       =     sd_pipe(**lowerCamelCase				,	negative_prompt=lowerCamelCase     )
        _UpperCamelCase   :				int       =     output.images
        _UpperCamelCase   :				str       =     image[0, -3:, -3:, -1]

        assert image.shape == (1, 6_4, 6_4, 3)

        _UpperCamelCase   :				Dict       =     np.array([0.4_7_7_1, 0.5_4_2_0, 0.4_6_8_3, 0.4_9_1_8, 0.5_6_3_6, 0.4_7_2_5, 0.5_2_3_0, 0.4_9_2_3, 0.5_0_1_5]     )

        assert np.abs(image_slice.flatten() - expected_slice     ).max() < 1e-2







@slow
@require_torch_gpu
class       __lowerCAmelCase     (				unittest.TestCase  ):



 '''simple docstring'''

 def     __snake_case    (	self				:						Optional[Any]     )       ->							Optional[int]:
        super().tearDown()
        gc.collect()
        torch.cuda.empty_cache()

 def     __snake_case    (	self				:						Dict				,	lowerCamelCase				:						Any=0     )       ->							Dict:
        _UpperCamelCase   :				List[str]       =     torch.Generator(device=lowerCamelCase     ).manual_seed(lowerCamelCase     )
        _UpperCamelCase   :				Tuple       =     {
            'prompt': 'a photograph of an astronaut riding a horse',
            'generator': generator,
            'num_inference_steps': 1_0,
            'guidance_scale': 7.5,
            'output_type': 'numpy',
            'parallel': 3,
            'debug': True,
        }
        return inputs

 def     __snake_case    (	self				:						List[Any]     )       ->							str:
        _UpperCamelCase   :				List[str]       =     'stabilityai/stable-diffusion-2-base'
        _UpperCamelCase   :				int       =     DDIMParallelScheduler.from_pretrained(lowerCamelCase				,	subfolder='scheduler'     )
        _UpperCamelCase   :				List[Any]       =     StableDiffusionParadigmsPipeline.from_pretrained(lowerCamelCase				,	scheduler=lowerCamelCase				,	safety_checker=lowerCamelCase     )
        pipe.to(lowerCamelCase     )
        pipe.set_progress_bar_config(disable=lowerCamelCase     )
        pipe.enable_attention_slicing()

        _UpperCamelCase   :				Optional[Any]       =     self.get_inputs()
        _UpperCamelCase   :				Optional[int]       =     pipe(**lowerCamelCase     ).images
        _UpperCamelCase   :				Optional[int]       =     image[0, -3:, -3:, -1].flatten()

        assert image.shape == (1, 5_1_2, 5_1_2, 3)

        _UpperCamelCase   :				int       =     np.array([0.9_6_2_2, 0.9_6_0_2, 0.9_7_4_8, 0.9_5_9_1, 0.9_6_3_0, 0.9_6_9_1, 0.9_6_6_1, 0.9_6_3_1, 0.9_7_4_1]     )

        assert np.abs(expected_slice - image_slice     ).max() < 1e-2



