import random


def downsample(df1, n_samples, do_segment_split=True, random_state=None):
    """
    Returns a downsampled version of *df1* so that it contains at most a ratio of *n_sample* samples.
    :param df1: The dataframe which should be downsampled.
    :param n_samples: The number of samples to include in the sample.
    :param do_segment_split: Whether the downsampling should be done per segment.
    :param random_state: Seed
    :return: A sample of df1 containing *n_samples* number of samples.
    """

    if random_state is not None:
        random.seed(random_state)

    if do_segment_split:
        df1_segments = list(sorted(df1.index.get_level_values('segment').unique()))
        samples_per_segment = len(df1)/len(df1_segments)

        n_segment_samples = int(n_samples / samples_per_segment)
        if n_segment_samples < len(df1_segments):
            sample_segments = random.sample(df1_segments, n_segment_samples)

            return df1.loc[sample_segments]
        else:
            return df1

    else:
        if n_samples < len(df1):
            print('N_samples: {}'.format(n_samples))
            sample_indices = random.sample(range(len(df1)), n_samples)
            return df1.iloc[sample_indices]
        else:
            print('N_samples: {}'.format(len(df1)))
            return df1