diff --git a/src/index.js b/src/index.js
index be505a6..a7c5b0c 100644
--- a/src/index.js
+++ b/src/index.js
@@ -5,13 +5,26 @@ import './index.css';
 import App from './App';
 import registerServiceWorker from './registerServiceWorker';
 
-const baseUrl = 'http://localhost:8080/dhis';
-// const baseUrl = 'https://idsr-afi.pheoc.go.ug/dhis';
+// const baseUrl = 'http://localhost:8080/dhis';
+// const baseUrl = 'https://cbs-academy.dhis2.org/demo';
+
+let baseUrl = '';
+if (process.env.NODE_ENV === 'development') {
+    baseUrl = 'http://localhost:8080/dhis/api/29';
+} else {
+    let urlArray = window.location.pathname.split('/');
+    let apiIndex = urlArray.indexOf('api');
+    if (apiIndex > 1) {
+        baseUrl = '/' + urlArray[apiIndex - 1] + '/api/29';
+    } else {
+        baseUrl = '/api/29';
+    }
+}
 
 init({
-    baseUrl: baseUrl + '/api/29'
+    baseUrl
 }).then(d2 => {
     ReactDOM.render(
-        <App d2={d2}/>, document.getElementById('root'));
+        <App d2={d2} baseUrl={baseUrl}/>, document.getElementById('root'));
     registerServiceWorker();
 }).catch(e => console.error);
diff --git a/src/programs/program.js b/src/programs/program.js
index 7346972..85bf419 100644
--- a/src/programs/program.js
+++ b/src/programs/program.js
@@ -16,11 +16,7 @@ import ExpandMoreIcon from '@material-ui/icons/ExpandMore';
 import ExportTable from './exportTable';
 
 import PropTypes from 'prop-types';
-
-
 // import Table from '@dhis2/d2-ui-table';
-
-
 const styles = theme => ({
     root: {
         flexGrow: 1,
@@ -102,18 +98,97 @@ class Program extends React.Component {
     }
 
     render() {
-        const {classes, d2} = this.props;
+        const {classes} = this.props;
         let progress = '';
+        let selection = '';
+
         if (this.integrationStore.loading) {
             progress = <LinearProgress variant="indeterminate"/>;
         }
+
+        if (this.integrationStore.program !== null) {
+            let current = '';
+            if (this.integrationStore.selectedStage !== null) {
+                current = <div>
+                    <InputField
+                        id="filter"
+                        label="Filter"
+                        type="text"
+                        fullWidth
+                        value={this.integrationStore.dataElementFilterText}
+                        onChange={(value) => this.integrationStore.dataElementFilterChange(value)}
+                    />
+
+                    <GroupEditor
+                        itemStore={this.integrationStore.dataElementStore}
+                        assignedItemStore={this.integrationStore.currentDataElementStore}
+                        onAssignItems={this.integrationStore.assignDataElements}
+                        onRemoveItems={this.integrationStore.unAssignDataElements}
+                        height={150}
+                        filterText={this.integrationStore.dataElementFilterText}
+                    />
+                </div>;
+            }
+            selection = <Tabs>
+                <Tab label='Attributes'>
+                    <InputField
+                        id="filter"
+                        label="Filter"
+                        type="text"
+                        fullWidth
+                        value={this.integrationStore.filterText}
+                        onChange={(value) => this.integrationStore.filterChange(value)}
+                    />
+                    <GroupEditor
+                        itemStore={this.integrationStore.itemStore}
+                        assignedItemStore={this.integrationStore.assignedItemStore}
+                        onAssignItems={this.integrationStore.assignItems}
+                        onRemoveItems={this.integrationStore.unAssignItems}
+                        height={150}
+                        filterText={this.integrationStore.filterText}
+                    />
+                </Tab>
+                <Tab label='Data Elements'>
+                    <br/>
+                    <Select
+                        placeholder="Select stage"
+                        value={this.integrationStore.selectedStage}
+                        options={this.integrationStore.programStages}
+                        labelKey="displayName"
+                        onChange={this.integrationStore.handleStageChange}
+                    />
+
+                    {current}
+
+                </Tab>
+
+                <Tab label='Other Columns'>
+                    <InputField
+                        id="filter"
+                        label="Filter"
+                        type="text"
+                        fullWidth
+                        value={this.integrationStore.filterOtherText}
+                        onChange={(value) => this.integrationStore.otherFilterChange(value)}
+                    />
+                    <GroupEditor
+                        itemStore={this.integrationStore.otherColumnStore}
+                        assignedItemStore={this.integrationStore.assignedOtherColumnStore}
+                        onAssignItems={this.integrationStore.assignOtherColumns}
+                        onRemoveItems={this.integrationStore.unAssignOtherColumns}
+                        height={150}
+                        filterText={this.integrationStore.filterOtherText}
+                    />
+                </Tab>
+            </Tabs>;
+        }
         return (
             <div>
                 {progress}
                 <div style={stylesx.div}>
                     <div zdepth={3} style={stylesx.paperLeft}>
                         <Select
-                            placeholder="Select one"
+                            placeholder="Select program"
                             value={this.integrationStore.program}
                             options={this.integrationStore.programs}
                             labelKey="displayName"
@@ -128,73 +203,7 @@ class Program extends React.Component {
                                     Attributes</Typography>
                             </ExpansionPanelSummary>
                             <ExpansionPanelDetails className={classes.block}>
-                                <Tabs>
-                                    <Tab label='Attributes'>
-                                        <InputField
-                                            id="filter"
-                                            label="Filter"
-                                            type="text"
-                                            fullWidth
-                                            value={this.integrationStore.filterText}
-                                            onChange={(value) => this.integrationStore.filterChange(value)}
-                                        />
-                                        <GroupEditor
-                                            itemStore={this.integrationStore.itemStore}
-                                            assignedItemStore={this.integrationStore.assignedItemStore}
-                                            onAssignItems={this.integrationStore.assignItems}
-                                            onRemoveItems={this.integrationStore.unAssignItems}
-                                            height={150}
-                                            filterText={this.integrationStore.filterText}
-                                        />
-                                    </Tab>
-                                    <Tab label='Data Elements'>
-                                        <br/>
-                                        <Select
-                                            placeholder="Select one"
-                                            value={this.integrationStore.selectedStage}
-                                            options={this.integrationStore.programStages}
-                                            labelKey="displayName"
-                                            onChange={this.integrationStore.handleStageChange}
-                                        />
-
-                                        <InputField
-                                            id="filter"
-                                            label="Filter"
-                                            type="text"
-                                            fullWidth
-                                            value={this.integrationStore.dataElementFilterText}
-                                            onChange={(value) => this.integrationStore.dataElementFilterChange(value)}
-                                        />
-
-                                        <GroupEditor
-                                            itemStore={this.integrationStore.dataElementStore}
-                                            assignedItemStore={this.integrationStore.assignedDataElementStore}
-                                            onAssignItems={this.integrationStore.assignDataElements}
-                                            onRemoveItems={this.integrationStore.unAssignDataElements}
-                                            height={150}
-                                            filterText={this.integrationStore.dataElementFilterText}
-                                        />
-                                    </Tab>
-
-                                    <Tab label='Other Columns'>
-                                        <InputField
-                                            id="filter"
-                                            label="Filter"
-                                            type="text"
-                                            fullWidth
-                                            value={this.integrationStore.filterOtherText}
-                                            onChange={(value) => this.integrationStore.otherFilterChange(value)}
-                                        />
-                                        <GroupEditor
-                                            itemStore={this.integrationStore.otherColumnStore}
-                                            assignedItemStore={this.integrationStore.assignedOtherColumnStore}
-                                            onAssignItems={this.integrationStore.assignOtherColumns}
-                                            onRemoveItems={this.integrationStore.unAssignOtherColumns}
-                                            height={150}
-                                            filterText={this.integrationStore.filterOtherText}
-                                        />
-                                    </Tab>
-                                </Tabs>
+                                {selection}
                                 <br/>
                             </ExpansionPanelDetails>
                         </ExpansionPanel>
@@ -246,7 +255,7 @@ class Program extends React.Component {
                             variant="contained"
                             color="primary"
                             onClick={this.integrationStore.downloadData}>
-                            Download
+                            {this.integrationStore.downloadLabel}
                         </Button>
                         <ExportTable/>
                     </div>
diff --git a/src/stores/IntegrationStore.js b/src/stores/IntegrationStore.js
index 8a2ce1a..51a8e61 100644
--- a/src/stores/IntegrationStore.js
+++ b/src/stores/IntegrationStore.js
@@ -41,7 +41,7 @@ class IntegrationStore {
     @observable assignedItemStore = Store.create();
 
     @observable dataElementStore = Store.create();
-    @observable assignedDataElementStore = Store.create();
+    @observable assignedDataElementStore = {};
 
     @observable otherColumnStore = Store.create();
     @observable assignedOtherColumnStore = Store.create();
@@ -106,8 +106,9 @@ class IntegrationStore {
     @observable expanded;
     @observable period;
     @observable selectedAttributes = [];
-    @observable selectedElements = [];
+    @observable selectedElements = {};
     @observable selectedOthers = [];
+    @observable downloadLabel = 'Download'
 
     constructor() {
         this.itemStore.state = [];
@@ -149,6 +150,18 @@ class IntegrationStore {
             program.programStages = program.programStages.toArray();
             program.organisationUnits = program.organisationUnits.toArray();
 
+            let stores = {};
+            let elements = {};
+
+            program.programStages.forEach(stage => {
+                const store = Store.create();
+                store.state = [];
+                stores = {...stores, ..._.fromPairs([[stage.id, store]])};
+                elements = {...elements, ..._.fromPairs([[stage.id, []]])}
+            });
+            this.selectedElements = elements;
+            this.assignedDataElementStore = stores;
+
             const items = program.programTrackedEntityAttributes.map(attribute => {
                 return {text: attribute.displayName, value: attribute.id};
             });
@@ -324,18 +337,18 @@ class IntegrationStore {
     };
 
     @action assignDataElements = (items) => {
-        const assigned = this.assignedDataElementStore.state.concat(items);
-        this.assignedDataElementStore.setState(assigned);
-        this.selectedElements = assigned;
+        const assigned = this.assignedDataElementStore[this.selectedStage.id].state.concat(items);
+        this.assignedDataElementStore[this.selectedStage.id].setState(assigned);
+        this.selectedElements[this.selectedStage.id] = assigned;
         return Promise.resolve();
     };
 
     @action unAssignDataElements = (items) => {
-        const assigned = this.assignedDataElementStore
+        const assigned = this.assignedDataElementStore[this.selectedStage.id]
             .state
             .filter(item => items.indexOf(item) === -1);
-        this.assignedDataElementStore.setState(assigned);
-        this.selectedElements = assigned;
+        this.assignedDataElementStore[this.selectedStage.id].setState(assigned);
+        this.selectedElements[this.selectedStage.id] = assigned;
         return Promise.resolve();
     };
 
@@ -426,7 +439,7 @@ class IntegrationStore {
         let view = this.sqlViews['analytics'];
         this.rows = [];
         this.headers = [];
-        this.rowsPerPage = 15;
+        this.rowsPerPage = 10;
         this.page = 0;
         if (view !== null && view !== undefined && view !== '') {
             api.get('sqlViews/' + view).then(action(sqlView => {
@@ -461,8 +474,9 @@ class IntegrationStore {
 
     @action downloadData = () => {
         const rowSize = (new TextEncoder('utf-8').encode(this.rows[0])).length;
-        const rowsToFetchAndDownload = Math.floor((50 * 1024 * 1024) / (4 * rowSize));
+        const rowsToFetchAndDownload = Math.floor((30 * 1024 * 1024) / (4 * rowSize));
         const numberOfPages = Math.ceil(this.total / rowsToFetchAndDownload);
+        this.downloadLabel = 'Downloading ' + numberOfPages + ' files';
         for (let i = 1; i <= numberOfPages; i++) {
             this.fetch(this.sqlViews['analytics'], i, rowsToFetchAndDownload);
         }
@@ -518,25 +532,36 @@ class IntegrationStore {
 
         const otherColumnString = otherColumns.join(',');
         let attributeString = '';
-        let dataElementString = '';
+        let elementsArray = [];
 
-        if (this.isTracker) {
+        if (this.isTracker && this.selectedStage !== null) {
             table = 'analytics_enrollment_' + this.program.id;
 
             attributeString = attributes.map(a => {
                 return '"' + a + '"';
             }).join(',');
 
-            dataElementString = dataElements.map(element => {
-                return '"' + this.selectedStage.id + '_' + element + '"';
-            }).join(',');
+            _.forOwn(dataElements, (e, s) => {
+                if (e.length > 0) {
+                    const r = e.map(element => {
+                        return '"' + s + '_' + element + '"';
+                    }).join(',');
+                    elementsArray = [...elementsArray, r];
+                }
+            });
 
 
         } else if (this.program !== null) {
             table = 'analytics_event_' + this.program.id;
-            dataElementString = dataElements.map(e => {
-                return '"' + e + '"';
-            }).join(',');
+
+            _.forOwn(dataElements, e => {
+                if (e.length > 0) {
+                    const r = e.map(el => {
+                        return '"' + el + '"';
+                    }).join(',');
+                    elementsArray = [...elementsArray, r];
+                }
+            });
         }
 
         if (attributeString !== null && attributeString !== '') {
@@ -547,8 +572,8 @@ class IntegrationStore {
             finalColumns = [...finalColumns, otherColumnString];
         }
 
-        if (dataElementString !== null && dataElementString !== '') {
-            finalColumns = [...finalColumns, dataElementString];
+        if (elementsArray.length > 0) {
+            finalColumns = [...finalColumns, elementsArray.join(',')];
         }
 
         if (finalColumns.length > 0) {
@@ -562,6 +587,10 @@ class IntegrationStore {
         return this.rows.length > 0;
     }
 
+    @computed get currentDataElementStore() {
+        return this.assignedDataElementStore[this.selectedStage.id];
+    }
+
     @computed get canUpdate() {
         return this.program !== null && this.query.columns !== '' && this.query.table !== '';
     }