<commit_before>define([
    'kb_lib/jsonRpc/genericClient',
    'kb_service/utils',
    'json!../data/groups.json',
    'json!../data/schemas.json',
    '../constants'
], (
    GenericClient,
    {objectInfoToObject},
    groupsData,
    schemasData,

    {MAX_SAMPLES}
) => {

    const schemas = schemasData.reduce((schemas, schema) => {
        schemas[schema.kbase.sample.key] = schema;
        return schemas;
    }, {});
    class SampleSetModel {
        constructor({runtime, workspaceId, objectId, objectVersion}) {
            this.runtime = runtime;
            this.workspaceId = workspaceId;
            this.objectId = objectId;
            this.objectVersion = objectVersion;
            this.sampleSetRef = `${workspaceId}/${objectId}/${objectVersion}`;

            this.cache = {
                objectInfos: {},
                objects: {},
                userProfiles: {},
                samples: {},
                dataLinksBySample: {}
            };
        }

        // Core data fetchers, which also cache.
        // FWIW we should try to clamp down on caching, as it can become
        // stale, and I'm not sure users would necessarily be aware
        // of this.

        // TODO: add caching!

        async getJSON(filePath) {
            const currentPath = module.uri.split('/').slice(1, -1).join('/');
            const response = await fetch(`${currentPath}/${filePath}.json`, {
                method: 'GET'
            });
            if (response.status !== 200) {
                throw new Error(`Cannot load JSON file ${filePath}`);
            }
            try {
                return await response.json();
            } catch (ex) {
                console.error('ERROR parsing JSON file', ex);
                throw new Error(`Error parsing JSON file "${filePath}: ${ex.message}`);
            }
        }

        /**
         *
         * @param fieldKeys {Array[string]}
         * @returns {string[]}
         */
        async getFieldSchemas(fieldKeys) {
            return Promise.all(fieldKeys.map((fieldKey) => {
                return schemas[fieldKey];
            }));
        }

        /**
         *
         * @returns {Promise<any>}
         */
        async getFieldGroups() {
            return groupsData;
        }

        async getObjectInfos(upas) {
            const upasToGet = [];
            const objectInfos = {};

            // First extract any in the cache.
            for (const upa of upas) {
                if (upa in this.cache.objectInfos) {
                    objectInfos[upa] = this.cache.objectInfos[upa];
                } else {
                    upasToGet.push(upa);
                }
            }

            // If there are any left to get, get them.
            if (upasToGet.length > 0) {
                const params = {
                    objects: upasToGet.map((ref) => {
                        return {
                            ref
                        };
                    }),
                    includeMetadata: 1,
                    ignoreErrors: 1
                };
                const workspaceClient = new GenericClient({
                    module: 'Workspace',
                    url: this.runtime.config('services.workspace.url'),
                    token: this.runtime.service('session').getAuthToken()
                });

                const [{infos}] = await workspaceClient.callFunc('get_object_info3', [params]);
                infos
                    .filter((info) => {
                        return info !== null;
                    })
                    .forEach((info) => {
                        const asObject = objectInfoToObject(info);
                        this.cache.objectInfos[asObject.ref] = asObject;
                        objectInfos[asObject.ref] = asObject;
                    });
            }

            return objectInfos;
        }

        sampleToRef({id, version}) {
            return `${id}/${version}`;
        }

        async getDataLinks(samplesToGet) {
            // NB links are cached by sample ref.
            const sampleLinksToFetch = samplesToGet.filter((sample) => {
                const key = this.sampleToRef(sample);
                return !(key in this.cache.dataLinksBySample);
            });

            if (sampleLinksToFetch.length > 0) {
                const sampleService = new GenericClient({
                    module: 'SampleService',
                    url: this.runtime.config('services.SampleService.url'),
                    token: this.runtime.service('session').getAuthToken()
                });

                // const start = Date.now();
                // const links = await Promise.all(sampleLinksToFetch.map(async ({id, version}) => {
                //     const [{links}] = await sampleService
                //         .callFunc('get_data_links_from_sample', [{id, version}]);
                //     return {links, id, version};

                // }));


                const result = await sampleService
                    .callFunc('get_data_links_from_sample_set', [{
                        sample_ids: sampleLinksToFetch,
                        effective_time: Date.now()
                    }]);

                const [{links}] = result;

                links.forEach((link) => {
                    // const {id, version} = sampleLinksToFetch[index];
                    const {id, version} = link;
                    const key = this.sampleToRef({id, version});
                    if (!(key in this.cache.dataLinksBySample)) {
                        this.cache.dataLinksBySample[key] = [];
                    }
                    this.cache.dataLinksBySample[key].push(link);
                });

                // for (const {id, version} of sampleLinksToFetch) {
                //     const [{links}] = await sampleService
                //         .callFunc('get_data_links_from_sample', [{id, version}]);
                //     const key = this.sampleToRef({id, version});
                //     this.cache.dataLinksBySample[key] = links;
                // }
            }
            return samplesToGet.map((sample) => {
                const key = this.sampleToRef(sample);
                return this.cache.dataLinksBySample[key];
            });
        }

        async getSamples(samplesToGet) {
            const samplesToFetch = samplesToGet.filter(({id, version}) => {
                const key = `${id}.${version}`;
                return !(key in this.cache.samples);
            })
                .map(({id, version}) => {
                    return {id, version};
                });

            if (samplesToFetch.length > 0) {
                const sampleService = new GenericClient({
                    module: 'SampleService',
                    url: this.runtime.config('services.SampleService.url'),
                    token: this.runtime.service('session').getAuthToken()
                });

                const [samples] = await sampleService
                    .callFunc('get_samples', [{
                        samples: samplesToFetch
                    }]);

                for (const sample of samples) {
                    const {id, version} = sample;
                    const key = `${id}.${version}`;
                    this.cache.samples[key] = sample;
                }
            }

            return samplesToGet.map(({id, version}) => {
                const key = `${id}.${version}`;
                return this.cache.samples[key];
            });
        }

        async getSampleSet() {
            if (this.sampleSetRef in this.cache.objects) {
                return this.cache.objects[this.sampleSetRef];
            }

            const workspaceClient = new GenericClient({
                module: 'Workspace',
                url: this.runtime.config('services.workspace.url'),
                token: this.runtime.service('session').getAuthToken()
            });

            const [result] = await workspaceClient.callFunc('get_objects2', [{
                objects: [
                    {
                        ref: this.sampleSetRef
                    }]
            }]);

            if (!result.data[0]) {
                throw new Error('SampleSet not found');
            }
            const {info, data} = result.data[0];
            const infoAsObject = objectInfoToObject(info);
            this.cache.objects[infoAsObject.ref] = {info: infoAsObject, data};
            return {info: infoAsObject, data};
        }

        async getUserProfiles(usernames) {
            const usersToGet = usernames.filter((username) => {
                return !(username in this.cache.userProfiles);
            });

            if (usersToGet.length > 0) {
                const userProfileService = new GenericClient({
                    module: 'UserProfile',
                    url: this.runtime.config('services.UserProfile.url'),
                    token: this.runtime.service('session').getAuthToken()
                });

                const [profiles] = await userProfileService.callFunc('get_user_profile', [usernames]);
                for (const profile of profiles) {
                    this.cache.userProfiles[profile.user.username] = profile;
                }
            }
            return usernames.map((username) => {
                return this.cache.userProfiles[username];
            });
        }

        getKeysForSamples(samples) {
            const fieldKeys = new Set();
            samples.forEach((sample) => {
                Object.keys(sample.node_tree[0].meta_controlled).forEach((key) => {
                    fieldKeys.add(key);
                });
                Object.keys(sample.node_tree[0].meta_user).forEach((key) => {
                    fieldKeys.add(key);
                });
            });
            return Array.from(fieldKeys);
        }

        async getTheSheBang() {
            let start = Date.now();
            const start0 = start;
            /******
             * Get Sample Set
             ******/
            const sampleSet = await this.getSampleSet();

            console.log('shebang: 1: getSampleSet', Date.now() - start);
            start = Date.now();

            const sampleCount = sampleSet.data.samples.length;

            /******
             * Get Samples
             ******/
            // Ensure we don't exceed hard-capped limit
            // DISABLED for now - figure out best way to handle this
            // sampleSet.samples = sampleSet.samples.slice(0, MAX_SAMPLES);

            const samples = await this.getSamples(sampleSet.data.samples);
            if (samples.length === 0) {
                // TODO: see if this is even possible.
                throw new Error('no samples in this set');
            }

            console.log('shebang: 2: getSamples', Date.now() - start);
            start = Date.now();

            const fieldKeys = new Set();
            samples.forEach((sample) => {
                Object.keys(sample.node_tree[0].meta_controlled).forEach((key) => {
                    fieldKeys.add(key);
                });
                Object.keys(sample.node_tree[0].meta_user).forEach((key) => {
                    fieldKeys.add(key);
                });
            });

            // Reform into a map of id to sample.
            // TODO: should include version in key?
            const samplesMap = samples.reduce((samplesMap, sample) => {
                samplesMap[sample.id] = sample;
                return samplesMap;
            }, {});

            // Create an array of samples in the same order as in the sampleset.
            const orderedSamples = sampleSet.data.samples.map((sampleSetItem) => {
                return samplesMap[sampleSetItem.id];
            });

            /******
             * Get data links
             ******/

            // Get the data links
            const rawDataLinks = await this.getDataLinks(orderedSamples);
            const samplesNLinks = samples.map((sample, index) => {
                return {
                    sample,
                    links: rawDataLinks[index]
                };
            });

            console.log('shebang: 3: getDataLinks', Date.now() - start);
            start = Date.now();

            /******
             * Get object infos from data links
             ******/
            // Get the upas
            const upas = new Set();
            for (const links of rawDataLinks) {
                // extract the upas
                for (const {upa} of links) {
                    upas.add(upa);
                }
            }

            // Get object info for each upa<commit_msg>better handling of data link caching<commit_after>define([
    'kb_lib/jsonRpc/genericClient',
    'kb_service/utils',
    'json!../data/groups.json',
    'json!../data/schemas.json',
    '../constants'
], (
    GenericClient,
    {objectInfoToObject},
    groupsData,
    schemasData,

    {MAX_SAMPLES}
) => {

    const schemas = schemasData.reduce((schemas, schema) => {
        schemas[schema.kbase.sample.key] = schema;
        return schemas;
    }, {});
    class SampleSetModel {
        constructor({runtime, workspaceId, objectId, objectVersion}) {
            this.runtime = runtime;
            this.workspaceId = workspaceId;
            this.objectId = objectId;
            this.objectVersion = objectVersion;
            this.sampleSetRef = `${workspaceId}/${objectId}/${objectVersion}`;

            this.cache = {
                objectInfos: {},
                objects: {},
                userProfiles: {},
                samples: {},
                dataLinksBySample: {}
            };
        }

        // Core data fetchers, which also cache.
        // FWIW we should try to clamp down on caching, as it can become
        // stale, and I'm not sure users would necessarily be aware
        // of this.

        // TODO: add caching!

        async getJSON(filePath) {
            const currentPath = module.uri.split('/').slice(1, -1).join('/');
            const response = await fetch(`${currentPath}/${filePath}.json`, {
                method: 'GET'
            });
            if (response.status !== 200) {
                throw new Error(`Cannot load JSON file ${filePath}`);
            }
            try {
                return await response.json();
            } catch (ex) {
                console.error('ERROR parsing JSON file', ex);
                throw new Error(`Error parsing JSON file "${filePath}: ${ex.message}`);
            }
        }

        /**
         *
         * @param fieldKeys {Array[string]}
         * @returns {string[]}
         */
        async getFieldSchemas(fieldKeys) {
            return Promise.all(fieldKeys.map((fieldKey) => {
                return schemas[fieldKey];
            }));
        }

        /**
         *
         * @returns {Promise<any>}
         */
        async getFieldGroups() {
            return groupsData;
        }

        async getObjectInfos(upas) {
            const upasToGet = [];
            const objectInfos = {};

            // First extract any in the cache.
            for (const upa of upas) {
                if (upa in this.cache.objectInfos) {
                    objectInfos[upa] = this.cache.objectInfos[upa];
                } else {
                    upasToGet.push(upa);
                }
            }

            // If there are any left to get, get them.
            if (upasToGet.length > 0) {
                const params = {
                    objects: upasToGet.map((ref) => {
                        return {
                            ref
                        };
                    }),
                    includeMetadata: 1,
                    ignoreErrors: 1
                };
                const workspaceClient = new GenericClient({
                    module: 'Workspace',
                    url: this.runtime.config('services.workspace.url'),
                    token: this.runtime.service('session').getAuthToken()
                });

                const [{infos}] = await workspaceClient.callFunc('get_object_info3', [params]);
                infos
                    .filter((info) => {
                        return info !== null;
                    })
                    .forEach((info) => {
                        const asObject = objectInfoToObject(info);
                        this.cache.objectInfos[asObject.ref] = asObject;
                        objectInfos[asObject.ref] = asObject;
                    });
            }

            return objectInfos;
        }

        sampleToRef({id, version}) {
            return `${id}/${version}`;
        }

        async getDataLinks(samplesToGet) {
            // NB links are cached by sample ref.
            const sampleLinksToFetch = samplesToGet.filter((sample) => {
                const key = this.sampleToRef(sample);
                return !(key in this.cache.dataLinksBySample);
            })
                .map(({id, version}) => {
                    return {id, version};
                });

            // Seed the cache with empty link set for each sample.
            // The data links api returns a list of links for all requested
            // samples, not a list for each sample.

            if (sampleLinksToFetch.length > 0) {
                for (const sample of sampleLinksToFetch) {
                    this.cache.dataLinksBySample[this.sampleToRef(sample)] = [];
                }

                const sampleService = new GenericClient({
                    module: 'SampleService',
                    url: this.runtime.config('services.SampleService.url'),
                    token: this.runtime.service('session').getAuthToken()
                });

                const result = await sampleService
                    .callFunc('get_data_links_from_sample_set', [{
                        sample_ids: sampleLinksToFetch,
                        effective_time: Date.now()
                    }]);

                const [{links}] = result;

                links.forEach((link) => {
                    const {id, version} = link;
                    const key = this.sampleToRef({id, version});
                    this.cache.dataLinksBySample[key].push(link);
                });
            }
            return samplesToGet.map((sample) => {
                const key = this.sampleToRef(sample);
                return this.cache.dataLinksBySample[key] || [];
            });
        }

        async getSamples(samplesToGet) {
            const samplesToFetch = samplesToGet.filter(({id, version}) => {
                const key = `${id}.${version}`;
                return !(key in this.cache.samples);
            })
                .map(({id, version}) => {
                    return {id, version};
                });

            if (samplesToFetch.length > 0) {
                const sampleService = new GenericClient({
                    module: 'SampleService',
                    url: this.runtime.config('services.SampleService.url'),
                    token: this.runtime.service('session').getAuthToken()
                });

                const [samples] = await sampleService
                    .callFunc('get_samples', [{
                        samples: samplesToFetch
                    }]);

                for (const sample of samples) {
                    const {id, version} = sample;
                    const key = `${id}.${version}`;
                    this.cache.samples[key] = sample;
                }
            }

            return samplesToGet.map(({id, version}) => {
                const key = `${id}.${version}`;
                return this.cache.samples[key];
            });
        }

        async getSampleSet() {
            if (this.sampleSetRef in this.cache.objects) {
                return this.cache.objects[this.sampleSetRef];
            }

            const workspaceClient = new GenericClient({
                module: 'Workspace',
                url: this.runtime.config('services.workspace.url'),
                token: this.runtime.service('session').getAuthToken()
            });

            const [result] = await workspaceClient.callFunc('get_objects2', [{
                objects: [
                    {
                        ref: this.sampleSetRef
                    }]
            }]);

            if (!result.data[0]) {
                throw new Error('SampleSet not found');
            }
            const {info, data} = result.data[0];
            const infoAsObject = objectInfoToObject(info);
            this.cache.objects[infoAsObject.ref] = {info: infoAsObject, data};
            return {info: infoAsObject, data};
        }

        async getUserProfiles(usernames) {
            const usersToGet = usernames.filter((username) => {
                return !(username in this.cache.userProfiles);
            });

            if (usersToGet.length > 0) {
                const userProfileService = new GenericClient({
                    module: 'UserProfile',
                    url: this.runtime.config('services.UserProfile.url'),
                    token: this.runtime.service('session').getAuthToken()
                });

                const [profiles] = await userProfileService.callFunc('get_user_profile', [usernames]);
                for (const profile of profiles) {
                    this.cache.userProfiles[profile.user.username] = profile;
                }
            }
            return usernames.map((username) => {
                return this.cache.userProfiles[username];
            });
        }

        getKeysForSamples(samples) {
            const fieldKeys = new Set();
            samples.forEach((sample) => {
                Object.keys(sample.node_tree[0].meta_controlled).forEach((key) => {
                    fieldKeys.add(key);
                });
                Object.keys(sample.node_tree[0].meta_user).forEach((key) => {
                    fieldKeys.add(key);
                });
            });
            return Array.from(fieldKeys);
        }

        async getTheSheBang() {
            let start = Date.now();
            const start0 = start;
            /******
             * Get Sample Set
             ******/
            const sampleSet = await this.getSampleSet();

            console.log('shebang: 1: getSampleSet', Date.now() - start);
            start = Date.now();

            const sampleCount = sampleSet.data.samples.length;

            /******
             * Get Samples
             ******/
            // Ensure we don't exceed hard-capped limit
            // DISABLED for now - figure out best way to handle this
            // sampleSet.samples = sampleSet.samples.slice(0, MAX_SAMPLES);

            const samples = await this.getSamples(sampleSet.data.samples);
            if (samples.length === 0) {
                // TODO: see if this is even possible.
                throw new Error('no samples in this set');
            }

            console.log('shebang: 2: getSamples', Date.now() - start);
            start = Date.now();

            const fieldKeys = new Set();
            samples.forEach((sample) => {
                Object.keys(sample.node_tree[0].meta_controlled).forEach((key) => {
                    fieldKeys.add(key);
                });
                Object.keys(sample.node_tree[0].meta_user).forEach((key) => {
                    fieldKeys.add(key);
                });
            });

            // Reform into a map of id to sample.
            // TODO: should include version in key?
            const samplesMap = samples.reduce((samplesMap, sample) => {
                samplesMap[sample.id] = sample;
                return samplesMap;
            }, {});

            // Create an array of samples in the same order as in the sampleset.
            const orderedSamples = sampleSet.data.samples.map((sampleSetItem) => {
                return samplesMap[sampleSetItem.id];
            });

            /******
             * Get data links
             ******/

            // Get the data links
            let rawDataLinks;
            try {
                rawDataLinks = await this.getDataLinks(orderedSamples);
            } catch (ex) {
                console.error('Error fetching data links', ex);
                throw new Error(`Error fetching data links: ${  ex.message}`);
            }
            console.log('shebang: 3a: getDataLinks', rawDataLinks);
            const samplesNLinks = samples.map((sample, index) => {
                return {
                    sample,
                    links: rawDataLinks[index]
                };
            });

            console.log('shebang: 3b: getDataLinks', Date.now() - start);
            start = Date.now();

            /******
             * Get object infos from data links
             ******/
            // Get the upas
            const upas = new Set();
            for (const links of rawDataLinks) {
                // extract the upas
                for (const {upa} of links) {
                    upas.add(upa);
                }
            }

            // Get object info for each upa
{'fraction_non_alphanumeric': 0.11636090463387633, 'fraction_numerical': 0.0020164478886604458, 'mean_word_length': 1.2745503597122303, 'pattern_counts': {'":': 0, '<': 5, '<?xml version=': 0, '>': 61, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '6650356', 'n_tokens_mistral': 6428, 'n_tokens_neox': 6038, 'n_words': 1340}